From afca9f33ace4fe27f06a785d2c3fbd03f0f2fd79 Mon Sep 17 00:00:00 2001 From: Antoine Groudiev Date: Fri, 1 May 2026 13:36:21 +0200 Subject: [PATCH 1/2] feat(ksos): support newton-sos solver Co-authored-by: Copilot --- CHANGELOG.md | 1 + example.py | 4 ++-- ksos_env.yml | 1 + ksos_tools/solvers/ksos.py | 40 ++++++++++++++++++++++++++++++-------- 4 files changed, 36 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index b8ae2dd..6ae6a86 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/example.py b/example.py index b34993c..0d4f020 100644 --- a/example.py +++ b/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") diff --git a/ksos_env.yml b/ksos_env.yml index 7197543..1194ba5 100644 --- a/ksos_env.yml +++ b/ksos_env.yml @@ -15,4 +15,5 @@ dependencies: - pytest - pip: - mosek + - newton_sos - -e . diff --git a/ksos_tools/solvers/ksos.py b/ksos_tools/solvers/ksos.py index 07c36e7..36ee873 100644 --- a/ksos_tools/solvers/ksos.py +++ b/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).reshape(-1, 1)) + # 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).reshape(-1, 1)) # TODO: avoid reshaping by changing the Rust code + 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( From 610a1799d4a294cde00986118e40e5a993da98b6 Mon Sep 17 00:00:00 2001 From: Antoine Groudiev Date: Fri, 1 May 2026 13:39:28 +0200 Subject: [PATCH 2/2] refactor(ksos): remove 1-dim reshape --- ksos_tools/solvers/ksos.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ksos_tools/solvers/ksos.py b/ksos_tools/solvers/ksos.py index 36ee873..112dff5 100644 --- a/ksos_tools/solvers/ksos.py +++ b/ksos_tools/solvers/ksos.py @@ -172,7 +172,7 @@ def solve( 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).reshape(-1, 1)) + 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": @@ -194,7 +194,7 @@ def solve( break 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).reshape(-1, 1)) # TODO: avoid reshaping by changing the Rust code + 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: