pybind: expose the unpreconditioned internal-basis gradient

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

src/vmecpp/cpp/vmecpp/vmec/pybind11/pybind_vmec.cc

@@ -217,9 +217,22 @@ class VmecModel {
   // decision vector. This is the body of Vmec::UpdateForwardModel with
   // caller-supplied counters; for free-boundary runs the caller should react to
   // `need_restart` (a vacuum-activation restart).
-  void Evaluate(int iter1, int iter2) {
+  // When precondition is true (default) this runs the full forward model and
+  // leaves decomposed_f holding the preconditioned search direction, exactly as
+  // the native solver uses it. When false, the forward model returns at the
+  // INVARIANT_RESIDUALS checkpoint (vmec.cc line ~836), so decomposed_f holds
+  // the raw, unpreconditioned force: the gradient of VMEC's augmented
+  // Lagrangian with respect to the (decomposed) state, including the
+  // lambda-constraint components. That raw gradient is what gradient-based
+  // optimizers minimizing the MHD energy functional need; mhd_energy is already
+  // set earlier in update(), so it is valid at the checkpoint too.
+  void Evaluate(int iter1, int iter2, bool precondition = true) {
     bool need_restart = false;
     std::string error_message;
+    const vmecpp::VmecCheckpoint checkpoint =
+        precondition ? vmecpp::VmecCheckpoint::NONE
+                     : vmecpp::VmecCheckpoint::INVARIANT_RESIDUALS;
+    const int checkpoint_after = precondition ? INT_MAX : 0;
     // Clear the restart reason before evaluating the forward model, exactly as
     // Vmec::Evolve does at its start (vmec.cc): the forward model only *sets* a
     // reason (BAD_JACOBIAN when the Jacobian flips, HUGE_INITIAL_FORCES at
@@ -240,7 +253,7 @@ class VmecModel {
           *vmec_->decomposed_x_[0], *vmec_->physical_x_[0],
           *vmec_->decomposed_f_[0], *vmec_->physical_f_[0], need_restart,
           last_preconditioner_update_, last_full_update_nestor_, vmec_->fc_,
-          iter1, iter2, vmecpp::VmecCheckpoint::NONE, INT_MAX,
+          iter1, iter2, checkpoint, checkpoint_after,
           /*verbose=*/false);
       if (!s.ok()) {
         error_message = std::string(s.status().message());
@@ -1177,7 +1190,8 @@ PYBIND11_MODULE(_vmecpp, m) {
   py::class_<VmecModel>(m, "VmecModel")
       .def_static("create", &VmecModel::Create, py::arg("indata"),
                   py::arg("ns"), py::arg("initial_state") = std::nullopt)
-      .def("evaluate", &VmecModel::Evaluate, py::arg("iter1"), py::arg("iter2"))
+      .def("evaluate", &VmecModel::Evaluate, py::arg("iter1"), py::arg("iter2"),
+           py::arg("precondition") = true)
       .def_property_readonly("need_restart", &VmecModel::need_restart)
       .def("perform_time_step", &VmecModel::PerformTimeStep,
            py::arg("velocity_scale"), py::arg("conjugation_parameter"),

tests/test_internal_gradient.py

@@ -0,0 +1,73 @@
+# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH
+# <info@proximafusion.com>
+#
+# 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.
+"""
+
+from pathlib import Path
+
+import numpy as np
+import pytest
+
+try:
+    from vmecpp.cpp import _vmecpp
+except ImportError:  # allow running against a directly-built extension
+    import _vmecpp
+
+SOLOVEV = Path(__file__).resolve().parents[1] / "examples" / "data" / "solovev.json"
+
+
+def _model(ns: int = 11):
+    indata = _vmecpp.VmecINDATA.from_file(str(SOLOVEV))
+    return _vmecpp.VmecModel.create(indata, ns)
+
+
+def test_raw_force_differs_from_preconditioned():
+    m = _model()
+    m.evaluate(2, 2, True)
+    f_prec = np.asarray(m.get_forces(), float)
+    m.evaluate(2, 2, False)
+    f_raw = np.asarray(m.get_forces(), float)
+
+    assert np.all(np.isfinite(f_raw))
+    assert np.linalg.norm(f_raw) > 0.0
+    # The preconditioner is a non-trivial metric: the two vectors are different
+    # in direction, not just scale.
+    cos = np.dot(f_prec, f_raw) / (np.linalg.norm(f_prec) * np.linalg.norm(f_raw))
+    assert abs(cos) < 0.99
+
+
+def test_raw_force_vanishes_at_equilibrium():
+    m = _model()
+    m.evaluate(2, 2, False)
+    f_initial = np.linalg.norm(np.asarray(m.get_forces(), float))
+
+    m.solve()
+    m.evaluate(2, 2, False)
+    f_converged = np.linalg.norm(np.asarray(m.get_forces(), float))
+
+    # The augmented-functional gradient is the equilibrium residual: it drops by
+    # many orders of magnitude once the native solver has converged.
+    assert f_converged < 1e-6 * f_initial
+
+
+def test_cold_start_is_excluded():
+    # evaluate(1, 2) is the cold-start special case (forces initialised to 1.0);
+    # the raw-gradient path uses iter1 >= 2, where the force is well defined.
+    m = _model()
+    m.evaluate(2, 2, False)
+    assert np.all(np.isfinite(np.asarray(m.get_forces(), float)))
+
+
+if __name__ == "__main__":
+    raise SystemExit(pytest.main([__file__, "-v"]))

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: