pybind: expose VMEC preconditioner operator + preconditioned JFNK

.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"

.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

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(

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,30 +117,56 @@ 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)),
         model.mhd_energy,
     )
 
 
+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}"
         )
 

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)

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"]))

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

tests/test_preconditioner.py

@@ -0,0 +1,63 @@
+# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH
+# <info@proximafusion.com>
+#
+# 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)

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: