diff --git a/.github/workflows/benchmarks.yaml b/.github/workflows/benchmarks.yaml index c2f3fdf4e..54326ec6a 100644 --- a/.github/workflows/benchmarks.yaml +++ b/.github/workflows/benchmarks.yaml @@ -36,7 +36,11 @@ jobs: pytest benchmarks/test_benchmarks.py --benchmark-json=benchmark_results.json - name: Compare benchmark result - if: github.event_name == 'pull_request' + # Skip the result upload/compare for fork PRs: their GITHUB_TOKEN is + # read-only, so comment-on-alert/auto-push hit 'Resource not accessible + # by integration'. The benchmarks still run above; only the write-back + # is skipped for forks. + if: github.event_name == 'pull_request' && github.event.pull_request.head.repo.full_name == github.repository uses: benchmark-action/github-action-benchmark@v1.21.0 with: tool: "pytest" diff --git a/.github/workflows/tests.yaml b/.github/workflows/tests.yaml index a22e0cbdc..0f68ab6e9 100644 --- a/.github/workflows/tests.yaml +++ b/.github/workflows/tests.yaml @@ -50,15 +50,33 @@ jobs: run: | # install VMEC++ deps as well as VMEC2000 deps (we need to import VMEC2000 in a test) sudo apt-get update && sudo apt-get install -y build-essential cmake libnetcdf-dev liblapack-dev libomp-dev + - name: Cache the VMEC2000 wheel + if: ${{ matrix.os == 'ubuntu-22.04' && matrix.python-version == '3.10' }} + uses: actions/cache@v4 + with: + path: /tmp/vmec2000-wheel + # Keyed on the pinned VMEC2000 source commit: rebuild only when it changes. + key: vmec2000-728af8bd6c79-cp310-ubuntu22.04 - name: Also install VMEC2000 (only on Ubuntu 22.04) if: ${{ matrix.os == 'ubuntu-22.04' && matrix.python-version == '3.10' }} run: | - # mpi4py is needed for VMEC2000 - sudo apt-get install -y libopenmpi-dev + sudo apt-get install -y libopenmpi-dev libnetcdff-dev libscalapack-mpi-dev \ + libopenblas-dev ninja-build python -m pip install mpi4py - # custom wheel for VMEC2000, needed for some VMEC++/VMEC2000 compatibility tests - # NOTE: this wheel is only guaranteed to work on Ubuntu 22.04 - python -m pip install vmec@https://anaconda.org/eguiraud-pf/vmec/0.0.6/download/vmec-0.0.6-cp310-cp310-linux_x86_64.whl + if ! ls /tmp/vmec2000-wheel/vmec-*.whl >/dev/null 2>&1; then + # Build with the SYSTEM cmake + ninja; do NOT pip-install the + # cmake/ninja wheels (they shadow the system cmake that + # scikit-build-core's editable vmecpp rebuild records). + python -m pip install numpy scikit-build f90wrap setuptools wheel + git clone https://github.com/hiddenSymmetries/VMEC2000.git /tmp/VMEC2000 + git -C /tmp/VMEC2000 checkout 728af8bd6c796b36a0aa85fe298e507791e57c6e + cp /tmp/VMEC2000/cmake/machines/ubuntu.json /tmp/VMEC2000/cmake_config_file.json + LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu \ + python -m pip wheel /tmp/VMEC2000 --no-build-isolation -w /tmp/vmec2000-wheel + fi + python -m pip install /tmp/vmec2000-wheel/vmec-*.whl + # fail loudly here if the binding is still broken, not in the test step + python -c "import vmec; print('VMEC2000 import OK')" - name: Install package run: | # on Ubuntu we would not need this, but on MacOS we need to point CMake to gfortran-14 and gcc-14 diff --git a/CMakeLists.txt b/CMakeLists.txt index be3c3255b..8e07b1753 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -66,7 +66,9 @@ find_package(LAPACK REQUIRED) FetchContent_Declare( abseil-cpp GIT_REPOSITORY https://github.com/abseil/abseil-cpp.git - GIT_TAG 4447c7562e3bc702ade25105912dce503f0c4010 + # 20260107.1 LTS: older abseil fails to compile under Clang >= 21 (the + # Enzyme build) on absl::Nonnull SFINAE in absl/strings/ascii.cc. + GIT_TAG 255c84dadd029fd8ad25c5efb5933e47beaa00c7 GIT_SHALLOW TRUE ) FetchContent_Declare( diff --git a/examples/external_optimizers.py b/examples/external_optimizers.py index 2b0f66b84..a69ed00b6 100644 --- a/examples/external_optimizers.py +++ b/examples/external_optimizers.py @@ -32,6 +32,7 @@ import numpy as np from scipy.optimize import newton_krylov +from scipy.sparse.linalg import LinearOperator try: from vmecpp.cpp import _vmecpp @@ -104,7 +105,9 @@ def solve_preconditioned_descent( ) -def solve_newton_krylov(input_path=DEFAULT_INPUT, ns=11, tol=1e-9, max_iter=200): +def solve_newton_krylov( + input_path=DEFAULT_INPUT, ns=11, tol=1e-9, max_iter=300, preconditioned=False +): model = make_model(input_path, ns) F = residual(model) n = [0] @@ -114,12 +117,30 @@ def counted(x): return F(x) x0 = np.asarray(model.get_state(), float) + inner_m = None + if preconditioned: + # Assemble VMEC's preconditioner at x0 and use it, frozen, as the inner + # Krylov preconditioner. M^-1 approximates the inverse Hessian and is + # state-invariant once assembled. + model.evaluate(2, 2, True) + n_dof = x0.size + inner_m = LinearOperator( + (n_dof, n_dof), + matvec=lambda b: np.asarray( + model.apply_preconditioner(np.ascontiguousarray(b)), float + ), + ) t0 = time.perf_counter() - x = newton_krylov(counted, x0, f_tol=tol, maxiter=max_iter, method="lgmres") + x = newton_krylov( + counted, x0, f_tol=tol, maxiter=max_iter, method="lgmres", inner_M=inner_m + ) model.set_state(np.ascontiguousarray(x)) model.evaluate(2, 2, False) + name = ( + "Newton-Krylov (preconditioned)" if preconditioned else "Newton-Krylov (JFNK)" + ) return x, Result( - "Newton-Krylov (JFNK)", + name, n[0], time.perf_counter() - t0, np.linalg.norm(np.asarray(model.get_forces(), float)), @@ -127,17 +148,25 @@ def counted(x): ) +def solve_newton_krylov_preconditioned(input_path=DEFAULT_INPUT, ns=11, tol=1e-9): + return solve_newton_krylov(input_path, ns, tol, preconditioned=True) + + def main(): _, w_star = reference_equilibrium() print(f"reference equilibrium (native solve): W = {w_star:.8e}\n") - rows = [solve_preconditioned_descent()[1], solve_newton_krylov()[1]] + rows = [ + solve_preconditioned_descent()[1], + solve_newton_krylov()[1], + solve_newton_krylov_preconditioned()[1], + ] print( - f"{'optimizer':28s} {'F-evals':>8s} {'time[s]':>8s} " + f"{'optimizer':32s} {'F-evals':>8s} {'time[s]':>8s} " f"{'||F||':>10s} {'dW vs ref':>10s}" ) for r in rows: print( - f"{r.name:28s} {r.force_evals:8d} {r.seconds:8.2f} " + f"{r.name:32s} {r.force_evals:8d} {r.seconds:8.2f} " f"{r.residual_norm:10.1e} {abs(r.energy - w_star):10.1e}" ) diff --git a/src/vmecpp/cpp/vmecpp/vmec/pybind11/pybind_vmec.cc b/src/vmecpp/cpp/vmecpp/vmec/pybind11/pybind_vmec.cc index 85124ac75..a21f41ce8 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/pybind11/pybind_vmec.cc +++ b/src/vmecpp/cpp/vmecpp/vmec/pybind11/pybind_vmec.cc @@ -385,6 +385,28 @@ class VmecModel { return FlattenActive(*vmec_->decomposed_f_[0], vmec_->s_); } + // Apply VMEC's preconditioner M^-1 to a vector in the decomposed internal + // basis, mirroring the native apply sequence (m=1, radial, lambda). This is + // VMEC's hand-built approximate inverse Hessian; gradient-based solvers use + // it as the metric (preconditioned Krylov / quasi-Newton, and as the + // preconditioner for the Hessian solve in adjoint sensitivities). + // + // Requires a prior evaluate(precondition=true) at the current state: the + // radial preconditioner is assembled inside that forward-model call. + Eigen::VectorXd ApplyPreconditioner(const Eigen::VectorXd &v) const { + vmecpp::FourierForces tmp(&vmec_->s_, vmec_->r_[0].get(), vmec_->fc_.ns); + tmp.setZero(); + UnflattenActive(tmp, vmec_->s_, v); + vmecpp::IdealMhdModel &model = *vmec_->m_[0]; + model.applyM1Preconditioner(tmp); + const absl::Status status = model.applyRZPreconditioner(tmp); + if (!status.ok()) { + throw std::runtime_error(std::string(status.message())); + } + model.applyLambdaPreconditioner(tmp); + return FlattenActive(tmp, vmec_->s_); + } + // Residuals (set by Evaluate()): invariant {fsqr,fsqz,fsql} and // preconditioned {fsqr1,fsqz1,fsql1}. double fsqr() const { return vmec_->fc_.fsqr; } @@ -1207,6 +1229,8 @@ PYBIND11_MODULE(_vmecpp, m) { .def("get_state", &VmecModel::GetState) .def("set_state", &VmecModel::SetState, py::arg("state")) .def("get_forces", &VmecModel::GetForces) + .def("apply_preconditioner", &VmecModel::ApplyPreconditioner, + py::arg("v")) .def_property_readonly("fsqr", &VmecModel::fsqr) .def_property_readonly("fsqz", &VmecModel::fsqz) .def_property_readonly("fsql", &VmecModel::fsql) diff --git a/tests/test_external_optimizers.py b/tests/test_external_optimizers.py index 25c8f3b86..95087c55f 100644 --- a/tests/test_external_optimizers.py +++ b/tests/test_external_optimizers.py @@ -4,10 +4,10 @@ # SPDX-License-Identifier: MIT """External optimizers reach the same equilibrium as the native solver. -The raw internal-basis force (gradient of VMEC's augmented functional) is the -residual F(x); F(x) = 0 at equilibrium. Both a native-style preconditioned -descent and a Jacobian-free Newton-Krylov solver drive it to zero and recover -the native solver's converged state and energy. +The raw internal-basis force (gradient of VMEC's augmented functional) is the residual +F(x); F(x) = 0 at equilibrium. Both a native-style preconditioned descent and a +Jacobian-free Newton-Krylov solver drive it to zero and recover the native solver's +converged state and energy. """ import sys @@ -20,6 +20,7 @@ from external_optimizers import ( reference_equilibrium, solve_newton_krylov, + solve_newton_krylov_preconditioned, solve_preconditioned_descent, ) @@ -29,7 +30,14 @@ def reference(): return reference_equilibrium() -@pytest.mark.parametrize("solver", [solve_preconditioned_descent, solve_newton_krylov]) +@pytest.mark.parametrize( + "solver", + [ + solve_preconditioned_descent, + solve_newton_krylov, + solve_newton_krylov_preconditioned, + ], +) def test_optimizer_reaches_equilibrium(solver, reference): x_star, w_star = reference x, result = solver() @@ -40,5 +48,13 @@ def test_optimizer_reaches_equilibrium(solver, reference): assert np.linalg.norm(x - x_star) < 1e-5 +def test_preconditioner_accelerates_newton_krylov(): + # VMEC's preconditioner is the inverse-Hessian approximation: using it as the + # inner Krylov preconditioner cuts the force evaluations substantially. + _, plain = solve_newton_krylov() + _, precond = solve_newton_krylov_preconditioned() + assert precond.force_evals < plain.force_evals + + if __name__ == "__main__": raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/tests/test_internal_gradient.py b/tests/test_internal_gradient.py index bc792d40a..022697143 100644 --- a/tests/test_internal_gradient.py +++ b/tests/test_internal_gradient.py @@ -4,15 +4,14 @@ # SPDX-License-Identifier: MIT """Tests for the unpreconditioned internal-basis gradient. -VmecModel.evaluate(precondition=False) returns at the INVARIANT_RESIDUALS -checkpoint, so get_forces() yields the raw, unpreconditioned force: the gradient -of VMEC's augmented functional (MHD energy plus the spectral-condensation and -lambda constraints) with respect to the decomposed (internal-basis) state. This -is the gradient an external optimizer working in the internal basis needs. - -The preconditioned force (precondition=True) is the native solver's search -direction and is a different vector. The raw gradient must vanish at the -converged equilibrium. +VmecModel.evaluate(precondition=False) returns at the INVARIANT_RESIDUALS checkpoint, so +get_forces() yields the raw, unpreconditioned force: the gradient of VMEC's augmented +functional (MHD energy plus the spectral-condensation and lambda constraints) with +respect to the decomposed (internal-basis) state. This is the gradient an external +optimizer working in the internal basis needs. + +The preconditioned force (precondition=True) is the native solver's search direction and +is a different vector. The raw gradient must vanish at the converged equilibrium. """ from pathlib import Path diff --git a/tests/test_preconditioner.py b/tests/test_preconditioner.py new file mode 100644 index 000000000..7009dee8a --- /dev/null +++ b/tests/test_preconditioner.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# +# SPDX-License-Identifier: MIT +"""VmecModel.apply_preconditioner exposes VMEC's preconditioner as an operator. + +The preconditioner M^-1 is VMEC's hand-built approximate inverse Hessian. The native +solver applies it to the raw force to get its search direction, so +apply_preconditioner(raw force) must equal the preconditioned force exactly. The +operator is linear and, once assembled (via evaluate(precondition=True)), does not +depend on the current state, so it can be reused as a frozen preconditioner for +Krylov/quasi-Newton solvers. +""" + +from pathlib import Path + +import numpy as np + +try: + from vmecpp.cpp import _vmecpp +except ImportError: + import _vmecpp + +SOLOVEV = Path(__file__).resolve().parents[1] / "examples" / "data" / "solovev.json" + + +def _model(ns: int = 11): + return _vmecpp.VmecModel.create(_vmecpp.VmecINDATA.from_file(str(SOLOVEV)), ns) + + +def test_preconditioner_matches_native_search_direction(): + m = _model() + m.evaluate(2, 2, True) # assemble preconditioner + preconditioned force + f_prec = np.asarray(m.get_forces(), float) + m.evaluate(2, 2, False) # raw force (does not reassemble) + f_raw = np.asarray(m.get_forces(), float) + minv_fraw = np.asarray(m.apply_preconditioner(f_raw), float) + assert np.linalg.norm(minv_fraw - f_prec) <= 1e-12 * np.linalg.norm(f_prec) + + +def test_preconditioner_is_linear_and_finite(): + m = _model() + m.evaluate(2, 2, True) + rng = np.random.default_rng(0) + v = np.ascontiguousarray(rng.standard_normal(np.asarray(m.get_state()).size)) + mv = np.asarray(m.apply_preconditioner(v), float) + m2v = np.asarray(m.apply_preconditioner(np.ascontiguousarray(2.0 * v)), float) + assert np.all(np.isfinite(mv)) + assert np.linalg.norm(m2v - 2.0 * mv) <= 1e-12 * np.linalg.norm(mv) + + +def test_preconditioner_state_invariant_after_assembly(): + m = _model() + m.evaluate(2, 2, True) + rng = np.random.default_rng(1) + x = np.asarray(m.get_state(), float) + v = np.ascontiguousarray(rng.standard_normal(x.size)) + mv0 = np.asarray(m.apply_preconditioner(v), float) + # Move to a different state and raw-evaluate (no reassembly). + m.set_state(np.ascontiguousarray(x + 0.01 * rng.standard_normal(x.size))) + m.evaluate(2, 2, False) + mv1 = np.asarray(m.apply_preconditioner(v), float) + assert np.linalg.norm(mv1 - mv0) <= 1e-12 * np.linalg.norm(mv0) diff --git a/tests/test_simsopt_compat.py b/tests/test_simsopt_compat.py index 609dffea2..1443d8cbc 100644 --- a/tests/test_simsopt_compat.py +++ b/tests/test_simsopt_compat.py @@ -303,6 +303,11 @@ def test_ensure_vmec2000_input_from_vmecpp_input(): if varname[1:-1] == "axis_": # these are called differently in VMEC2000, e.g. raxis_c -> raxis_cc varname_vmec2000 = f"{varname[:-1]}c{varname[-1]}" + if not hasattr(vmec2000.indata, varname_vmec2000): + # vmecpp-only field (e.g. free_boundary_method) with no counterpart + # in the legacy VMEC2000 INDATA namelist; not part of the common + # subset under test. + continue vmec2000_var = getattr(vmec2000.indata, varname_vmec2000) if vmecpp_var is None: