Interface the new solver with the toolbox

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

@@ -94,7 +94,8 @@ 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
     - name: Test with pytest
       env:
           MOSEKLM_LICENSE_FILE: ${{ secrets.MOSEK_LICENSE }}

CHANGELOG.md

@@ -3,15 +3,19 @@
 This file is used to track changes made to the project over time.
 
 ## [Unreleased]
+- none
+
+## [0.0.3] - 2026-08-31
 ### 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
+## [0.0.2] - 2025-11-14
 ### Additions
 - Add penalty parameter for soft-constrained version used with polynomial kernels
 - Create helper functions for monomial feature function creation for polynomial kernel
 
 ### Improvements
 - Improve timing calculations
-- Improve efficiency of surrogate function calculation
+- Improve efficiency of surrogate function calculation

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,19 @@
+name: ksos
+channels:
+  - conda-forge
+  - defaults
+dependencies:
+  - python=3.12
+  - pip
+  - setuptools
+  - pytest
+
+  - cvxpy
+  - eigenpy
+  - matplotlib
+  - numpy
+  - scipy
+  - pytest
+  - pip:
+    - mosek
+    - -e .

ksos_tools/__init__.py

@@ -1 +1 @@
-__version__ = "0.0.2"
+__version__ = "0.0.3"

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,26 @@ def solve(
             else:
                 problem.register_fixed_samples(samples, f, None)
 
-            if solver != "naive":
+            if solver == "newton-rs":
+                import newton_sos
+
+                try:
+                    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)
+                except Exception as exc:
+                    warnings.warn(
+                        f"Warning: Error encountered in newton-rs: {exc}"
+                    )
+                    info["cost"] = None
+                    info["success"] = False
+                    info["status"] = f"Error in newton-rs: {exc}"
+                    return None, info
+            elif solver != "naive":
                 success = problem.initialize_kernel(
                     sigma, kernel, verbose=verbose, llt_method=llt_method
                 )
@@ -186,12 +206,35 @@ 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
+
+                    try:
+                        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)
+                    except Exception as exc:
+                        warnings.warn(
+                            f"Warning: Error encountered in newton-rs: {exc}"
+                        )
+                        fail_count += 1
+                        if fail_count >= MAX_FAIL_COUNT or sampling == "linspace":
+                            info["cost"] = None
+                            info["success"] = False
+                            info["status"] = f"Error in newton-rs: {exc}"
+                            return None, info
+                        continue
                     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 +266,33 @@ def solve(
                 verbose=verbose,
                 return_B=return_B,
             )
+        elif solver == "newton-rs":
+            try:
+                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,
+                }
+            except Exception as exc:  # pragma: no cover - depends on native solver
+                warnings.warn(
+                    f"Warning: newton-rs failed due to an ill-posed kernel: {exc}"
+                )
+                info["cost"] = None
+                info["success"] = False
+                info["status"] = "Kernel matrix not PSD"
+                return None, info
         elif solver == "newton-features":
             problem.use_K = False
             z, info_here = newton.damped_newton_advanced(

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "ksos_tools"
-version = "0.0.2"
+version = "0.0.3"
 authors = [
     { name = "Antoine Groudiev", email = "antoine.groudiev@ens.psl.eu" },
     { name = "Frederike Dümbgen", email = "frederike.duembgen@gmail.com" }

tests/test_benchmarks.py

@@ -1,9 +1,11 @@
+import warnings
 import matplotlib
 import matplotlib.colors
 import matplotlib.pylab as plt
 import numpy as np
 import pytest
 import os
+import importlib.util
 
 from ksos_tools.examples.benchmarks import ackley, rosenbrock, schwefel
 from ksos_tools.solvers import ksos
@@ -20,6 +22,15 @@
     sampling="linspace", return_all=False, return_B=False, verbose=False
 )
 
+if importlib.util.find_spec("newton_sos") is not None:
+    SOLVERS = ("MOSEK", "newton", "newton-rs", "newton-features", "newton-kernel")
+else:
+    SOLVERS = ("MOSEK", "newton", "newton-features", "newton-kernel")
+    warnings.warn(
+        "newton-rs solver is not available because newton_sos module is not installed. "
+        "Please install newton_sos to use this solver."
+    )
+
 
 def plot_solutions(center, radius, info, x_gt, f):
     # plot sample distributions vs. original cost
@@ -37,9 +48,7 @@ def plot_solutions(center, radius, info, x_gt, f):
     plt.show(block=False)
 
 
-@pytest.mark.parametrize(
-    "solver", ("MOSEK", "newton", "newton-features", "newton-kernel")
-)
+@pytest.mark.parametrize("solver", SOLVERS)
 @pytest.mark.parametrize("kernel", ("Laplace", "Gauss"))
 def test_ackley(solver, kernel, plot=False):
     f_here = lambda x: ackley(x)  # type: ignore # noqa: E731
@@ -72,7 +81,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,13 +134,11 @@ 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")
-)
+@pytest.mark.parametrize("solver", SOLVERS)
 @pytest.mark.parametrize("kernel", ("Laplace", "Gauss"))
 def test_rosenbrock(solver, kernel, plot=False):
     a = 1.0
@@ -167,7 +174,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
 
 
@@ -178,9 +185,7 @@ def test_rosenbrock(solver, kernel, plot=False):
     # test_rosenbrock(solver="newton-new", kernel="Gauss", plot=True)
     # 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"]
-    ):
+    for solver, kernel in itertools.product(SOLVERS, ["Gauss", "Laplace"]):
         print(f"========== testing {solver} {kernel} =============")
         print(f"----------         Ackley            -------------")
         test_ackley(solver=solver, kernel=kernel, plot=plot)

tests/test_errors.py

@@ -0,0 +1,42 @@
+import warnings
+
+import numpy as np
+import pytest
+
+from ksos_tools.solvers import ksos
+from tests.test_benchmarks import SOLVERS
+
+
+@pytest.mark.parametrize("solver", SOLVERS)
+def test_matrix_not_psd(solver):
+    # Test if the solver can handle a non-PSD kernel matrix properly.
+
+    samples = np.zeros((5, 1), dtype=float)
+    f = lambda x: float(np.sum(x**2))  # noqa: E731
+
+    with warnings.catch_warnings(record=True) as caught:
+        warnings.simplefilter("always")
+        z, info = ksos.solve(
+            f=f,
+            dim=1,
+            samples=samples,
+            solver=solver,
+            kernel="Gauss",
+            sigma=1.0,
+            epsilon=1e-6,
+            n_samples=5,
+            return_B=True,
+            verbose=False,
+        )
+
+    assert isinstance(info, dict)
+    assert z is None or np.all(np.isfinite(z))
+    assert len(caught) >= 0
+    if info.get("status") is not None:
+        assert info["status"] in {"Kernel matrix not PSD", "Solution extrapolated"}
+
+
+if __name__ == "__main__":
+    for solver in SOLVERS:
+        print(f"Testing solver: {solver}")
+        test_matrix_not_psd(solver)

tests/test_polynomial_problems.py

@@ -5,6 +5,8 @@
 import matplotlib.pylab as plt
 import numpy as np
 
+from tests.test_benchmarks import SOLVERS
+
 # Set matplotlib backend based on environment to avoid display issues
 if os.environ.get("DISPLAY", "") == "":
     # Headless environment (e.g., CI)
@@ -86,7 +88,7 @@ def test_newton_vs_mosek():
 
     z_dict = {}
     info_dict = {}
-    for solver in ["MOSEK", "newton", "newton-kernel", "newton-features"]:
+    for solver in SOLVERS:
         print(f"\n\nsolving with {solver}")
         t1 = time.time()
         z_solver, info_solver = ksos.solve(
@@ -115,9 +117,12 @@ 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 SOLVERS:
+        # compare to all solvers but MOSEK
+        if solver == "MOSEK":
+            continue
         # 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
@@ -207,8 +212,10 @@ def test_polynomial_kernel():
         if soft_constraints:
             solvers = ["MOSEK"]
         else:
-            # TODO: newton-features and newton-kernel not working. 
+            # 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: