Interface the new solver with the toolbox

.github/workflows/python-package-conda.yml

@@ -94,7 +94,10 @@ jobs:
         python-version: '3.12'
     - name: Install dependencies
       run: |
-        conda env create --file ksos_env.yml --name ksos_env
+        conda env create --file ksos_env_windows.yml --name ksos_env
+        conda activate ksos_env
+        python -m pip install --upgrade pip setuptools maturin
+        python -m pip install newton_sos
     - name: Test with pytest
       env:
           MOSEKLM_LICENSE_FILE: ${{ secrets.MOSEK_LICENSE }}

CHANGELOG.md

@@ -6,6 +6,7 @@ This file is used to track changes made to the project over time.
 ### Additions
 - Support for Sobol sequences in polynomial problem sampling
 - Support for periodic kernel
+- Support for [`newton-sos`](https://github.com/agroudiev/newton-sos) Rust-based solver
 
 ## [0.2.2] - 2025-11-14
 ### Additions

example.py

@@ -11,7 +11,7 @@
 # solver), newton-features or newton-kernel (more advanced solvers using
 # feature or kernel matrices, respectively; adding e.g. a linesearch option
 # and more advanced diagnostics to the original solver.
-solver = "newton"
+solver = "newton-rs"
 
 # Which kernel to use: use Gauss for smooth, Laplace for less smooth,
 # or provide a kernel of your choice.
@@ -31,5 +31,5 @@
     solver=solver,
     sigma=sigma,
 )
-print(f"Found solution: x={solution[0]:.4f}, f={info['cost']:.4f}")
+print(f"Found solution: x={solution[0].reshape(-1).item():.4f}, f={info['cost']:.4f}")
 print(f"True solution:  x={-np.pi/2:.4f}, f=-1")

ksos_env.yml

@@ -15,4 +15,5 @@ dependencies:
   - pytest
   - pip:
     - mosek
+    - newton_sos
     - -e .

ksos_env_windows.yml

@@ -0,0 +1,20 @@
+name: ksos
+channels:
+  - conda-forge
+  - defaults
+dependencies:
+  - python=3.12
+  - pip
+  - setuptools
+  - rust
+  - pytest
+
+  - cvxpy
+  - eigenpy
+  - matplotlib
+  - numpy
+  - scipy
+  - pytest
+  - pip:
+    - mosek
+    - -e .

ksos_tools/solvers/ksos.py

@@ -75,9 +75,9 @@ def solve(
     verbose: bool
         If True, prints the Sobolev norm and decay at each iteration.
     solver: str
-        The solver to use. Either 'newton', 'newton-original','MOSEK', 'SCS', or 'naive'.
+        The solver to use. Either 'newton', 'newton-rs','MOSEK', 'SCS', or 'naive'.
         - `newton`: uses the damped Newton method as suggested by Rudi et al.
-        - `newton-new`: uses a new interior-point Newton method.
+        - `newton-rs`: same as `newton`, using an efficient implementation in Rust
         - `naive`: retrieves the best sample.
         - others: uses CVXPY with the specified solver.
     max_iters_scs: int
@@ -126,6 +126,7 @@ def solve(
     assert isinstance(verbose, bool)
     assert solver in [
         "newton",
+        "newton-rs",
         "newton-features",
         "newton-kernel",
         "MOSEK",
@@ -169,7 +170,12 @@ def solve(
             else:
                 problem.register_fixed_samples(samples, f, None)
 
-            if solver != "naive":
+            if solver == "newton-rs":
+                import newton_sos
+                rs_problem = newton_sos.Problem(lambd, t, problem.samples.astype(np.float64), problem.f_samples.astype(np.float64))
+                # TODO: catch errors if any
+                rs_problem.initialize_native_kernel(kernel, sigma)
+            elif solver != "naive":
                 success = problem.initialize_kernel(
                     sigma, kernel, verbose=verbose, llt_method=llt_method
                 )
@@ -186,12 +192,17 @@ def solve(
                 )
                 if solver == "naive":
                     break
-
-                success = problem.initialize_kernel(
-                    sigma, kernel, verbose=verbose, llt_method=llt_method
-                )
-                if success:
+                elif solver == "newton-rs":
+                    import newton_sos
+                    rs_problem = newton_sos.Problem(lambd, t, problem.samples.astype(np.float64), problem.f_samples.astype(np.float64))
+                    rs_problem.initialize_native_kernel(kernel, sigma) # TODO: catch errors if any
                     break
+                else:
+                    success = problem.initialize_kernel(
+                        sigma, kernel, verbose=verbose, llt_method=llt_method
+                    )
+                    if success:
+                        break
 
                 fail_count += 1
                 if fail_count >= MAX_FAIL_COUNT or sampling == "linspace":
@@ -223,6 +234,19 @@ def solve(
                 verbose=verbose,
                 return_B=return_B,
             )
+        elif solver == "newton-rs":
+            solve_result = newton_sos.solve(rs_problem, max_iter=max_iters_newton, verbose=verbose, method="partial_piv_lu")
+            z = solve_result.z_hat
+            # TODO: lazy evaluation of phi and B
+            rs_problem.compute_phi()
+            info_here = {
+                "cost": solve_result.cost,
+                "alpha": solve_result.alpha,
+                "status": solve_result.status,
+                "success": solve_result.converged,
+                "B": solve_result.get_B(rs_problem),
+                # "X": X,
+            }
         elif solver == "newton-features":
             problem.use_K = False
             z, info_here = newton.damped_newton_advanced(

tests/test_benchmarks.py

@@ -38,7 +38,7 @@ def plot_solutions(center, radius, info, x_gt, f):
 
 
 @pytest.mark.parametrize(
-    "solver", ("MOSEK", "newton", "newton-features", "newton-kernel")
+    "solver", ("MOSEK", "newton", "newton-rs", "newton-features", "newton-kernel")
 )
 @pytest.mark.parametrize("kernel", ("Laplace", "Gauss"))
 def test_ackley(solver, kernel, plot=False):
@@ -72,7 +72,7 @@ def test_ackley(solver, kernel, plot=False):
     if plot:
         plot_solutions(center, radius, info, x_gt, f_here)
 
-    np.testing.assert_allclose(x_hat, x_gt, atol=0.5)
+    np.testing.assert_allclose(x_hat.flatten(), x_gt, atol=0.5)
     return
 
 
@@ -125,12 +125,12 @@ def f_here(x):
     if plot:
         plot_solutions(center, radius, info, x_gt, f_here)
 
-    np.testing.assert_allclose(x_hat, x_gt, rtol=1e-1)
+    np.testing.assert_allclose(x_hat.flatten(), x_gt, rtol=1e-1)
     return
 
 
 @pytest.mark.parametrize(
-    "solver", ("MOSEK", "newton", "newton-features", "newton-kernel")
+    "solver", ("MOSEK", "newton", "newton-rs", "newton-features", "newton-kernel")
 )
 @pytest.mark.parametrize("kernel", ("Laplace", "Gauss"))
 def test_rosenbrock(solver, kernel, plot=False):
@@ -167,7 +167,7 @@ def test_rosenbrock(solver, kernel, plot=False):
     if plot:
         plot_solutions(center, radius, info, x_gt, f_here)
 
-    np.testing.assert_allclose(x_hat, x_gt, atol=0.5)
+    np.testing.assert_allclose(x_hat.flatten(), x_gt, atol=0.5)
     return
 
 
@@ -179,7 +179,8 @@ def test_rosenbrock(solver, kernel, plot=False):
     # test_schwefel(solver="newton-new", kernel="Gauss", plot=True)
     plot = False
     for solver, kernel in itertools.product(
-        ["MOSEK", "newton", "newton-kernel", "newton-features"], ["Gauss", "Laplace"]
+        ["MOSEK", "newton", "newton-rs", "newton-kernel", "newton-features"], 
+        ["Gauss", "Laplace"]
     ):
         print(f"========== testing {solver} {kernel} =============")
         print(f"----------         Ackley            -------------")

tests/test_polynomial_problems.py

@@ -86,7 +86,7 @@ def test_newton_vs_mosek():
 
     z_dict = {}
     info_dict = {}
-    for solver in ["MOSEK", "newton", "newton-kernel", "newton-features"]:
+    for solver in ["MOSEK", "newton", "newton-rs", "newton-kernel", "newton-features"]:
         print(f"\n\nsolving with {solver}")
         t1 = time.time()
         z_solver, info_solver = ksos.solve(
@@ -115,9 +115,9 @@ def test_newton_vs_mosek():
         # assert alpha is feasible
         assert abs(np.sum(info_solver["alpha"]) - 1) <= 1e-10
 
-    for solver in ["newton", "newton-kernel", "newton-features"]:
+    for solver in ["newton", "newton-rs", "newton-kernel", "newton-features"]:
         # make sure both find the same solution
-        np.testing.assert_allclose(z_dict["MOSEK"], z_dict[solver], rtol=1e-2)
+        np.testing.assert_allclose(z_dict["MOSEK"], z_dict[solver].flatten(), rtol=1e-2)
         # make sure both find the same cost
         np.testing.assert_allclose(
             info_dict["MOSEK"]["cost"], info_dict[solver]["cost"], rtol=1e-2
@@ -209,6 +209,8 @@ def test_polynomial_kernel():
         else:
             # TODO: newton-features and newton-kernel not working. 
             # Should investigate why.
+            # newton-rs is intentionally left out cause it does not
+            # support polynomial kernels yet.
             solvers = ["MOSEK", "newton"]
 
         for solver in solvers: