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/src/vmecpp/cpp/vmecpp/vmec/pybind11/pybind_vmec.cc b/src/vmecpp/cpp/vmecpp/vmec/pybind11/pybind_vmec.cc index ccdf3f7af..85124ac75 100644 --- a/src/vmecpp/cpp/vmecpp/vmec/pybind11/pybind_vmec.cc +++ b/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_(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"), diff --git a/tests/test_internal_gradient.py b/tests/test_internal_gradient.py new file mode 100644 index 000000000..022697143 --- /dev/null +++ b/tests/test_internal_gradient.py @@ -0,0 +1,73 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# +# 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"])) 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: