diff --git a/.github/workflows/python-package-conda.yml b/.github/workflows/python-package-conda.yml index 82d5851..0103d4d 100644 --- a/.github/workflows/python-package-conda.yml +++ b/.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 }} 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_env_windows.yml b/ksos_env_windows.yml new file mode 100644 index 0000000..d929395 --- /dev/null +++ b/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 . diff --git a/ksos_tools/solvers/ksos.py b/ksos_tools/solvers/ksos.py index 07c36e7..112dff5 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)) + # 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( diff --git a/tests/test_benchmarks.py b/tests/test_benchmarks.py index f5c50fb..3d6b67b 100644 --- a/tests/test_benchmarks.py +++ b/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 -------------") diff --git a/tests/test_polynomial_problems.py b/tests/test_polynomial_problems.py index 119bca9..ce6afa6 100644 --- a/tests/test_polynomial_problems.py +++ b/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: