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 new file mode 100644 index 000000000..2b0f66b84 --- /dev/null +++ b/examples/external_optimizers.py @@ -0,0 +1,146 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# +# SPDX-License-Identifier: MIT +"""Drive a VMEC++ equilibrium from outside with general-purpose optimizers. + +VMEC's equilibrium is the stationary point of its augmented functional (MHD +energy plus the spectral-condensation and lambda constraints). The gradient of +that functional in the decomposed internal basis is the raw, unpreconditioned +force exposed by ``VmecModel.evaluate(precondition=False)`` (see +``tests/test_internal_gradient.py``). Finding the equilibrium is therefore the +root problem F(x) = 0, which any gradient/Hessian-based solver can attack while +reusing VMEC++'s forward model. + +This module wires that residual to two solvers and is used both by the benchmark +``main`` below and by ``tests/test_external_optimizers.py``: + +* native-style preconditioned descent (heavy-ball on the preconditioned search + direction, i.e. VMEC's own update), and +* Jacobian-free Newton-Krylov (matrix-free Hessian information). + +Both converge to the same equilibrium as the native solver. Quasi-Newton +root-finders without a preconditioner diverge on this stiff system, which is why +VMEC's preconditioner matters; exposing it as an operator is a follow-up. +""" + +from __future__ import annotations + +import time +from dataclasses import dataclass +from pathlib import Path + +import numpy as np +from scipy.optimize import newton_krylov + +try: + from vmecpp.cpp import _vmecpp +except ImportError: # directly-built extension on PYTHONPATH + import _vmecpp + +DEFAULT_INPUT = ( + Path(__file__).resolve().parents[1] / "examples" / "data" / "solovev.json" +) + + +def make_model(input_path: Path = DEFAULT_INPUT, ns: int = 11): + indata = _vmecpp.VmecINDATA.from_file(str(input_path)) + return _vmecpp.VmecModel.create(indata, ns) + + +def residual(model): + """Return F(x) = raw internal-basis force; F(x) = 0 at equilibrium.""" + + def F(x): + model.set_state(np.ascontiguousarray(x, dtype=float)) + model.evaluate(2, 2, False) + return np.asarray(model.get_forces(), dtype=float) + + return F + + +@dataclass +class Result: + name: str + force_evals: int + seconds: float + residual_norm: float + energy: float + + +def reference_equilibrium(input_path: Path = DEFAULT_INPUT, ns: int = 11): + model = make_model(input_path, ns) + model.solve() + model.evaluate(2, 2, False) + return np.asarray(model.get_state(), float), model.mhd_energy + + +def solve_preconditioned_descent( + input_path=DEFAULT_INPUT, ns=11, tol=1e-9, delt=0.9, momentum=0.5, max_iter=20000 +): + model = make_model(input_path, ns) + F = residual(model) + x = np.asarray(model.get_state(), float).copy() + v = np.zeros_like(x) + evals = 0 + t0 = time.perf_counter() + for _ in range(max_iter): + if np.linalg.norm(F(x)) < tol: + break + evals += 1 + model.set_state(np.ascontiguousarray(x)) + model.evaluate(2, 2, True) # preconditioned search direction + fprec = np.asarray(model.get_forces(), float) + v = momentum * v + delt * fprec + x = x + delt * v + model.set_state(np.ascontiguousarray(x)) + model.evaluate(2, 2, False) + return x, Result( + "preconditioned descent", + evals, + time.perf_counter() - t0, + np.linalg.norm(np.asarray(model.get_forces(), float)), + model.mhd_energy, + ) + + +def solve_newton_krylov(input_path=DEFAULT_INPUT, ns=11, tol=1e-9, max_iter=200): + model = make_model(input_path, ns) + F = residual(model) + n = [0] + + def counted(x): + n[0] += 1 + return F(x) + + x0 = np.asarray(model.get_state(), float) + t0 = time.perf_counter() + x = newton_krylov(counted, x0, f_tol=tol, maxiter=max_iter, method="lgmres") + model.set_state(np.ascontiguousarray(x)) + model.evaluate(2, 2, False) + return x, Result( + "Newton-Krylov (JFNK)", + n[0], + time.perf_counter() - t0, + np.linalg.norm(np.asarray(model.get_forces(), float)), + model.mhd_energy, + ) + + +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]] + print( + f"{'optimizer':28s} {'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.residual_norm:10.1e} {abs(r.energy - w_star):10.1e}" + ) + + +if __name__ == "__main__": + main() diff --git a/tests/test_external_optimizers.py b/tests/test_external_optimizers.py new file mode 100644 index 000000000..68b84562c --- /dev/null +++ b/tests/test_external_optimizers.py @@ -0,0 +1,44 @@ +# SPDX-FileCopyrightText: 2024-present Proxima Fusion GmbH +# +# +# 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. +""" + +import sys +from pathlib import Path + +import numpy as np +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "examples")) +from external_optimizers import ( + reference_equilibrium, + solve_newton_krylov, + solve_preconditioned_descent, +) + + +@pytest.fixture(scope="module") +def reference(): + return reference_equilibrium() + + +@pytest.mark.parametrize("solver", [solve_preconditioned_descent, solve_newton_krylov]) +def test_optimizer_reaches_equilibrium(solver, reference): + x_star, w_star = reference + x, result = solver() + # Force balance achieved. + assert result.residual_norm < 1e-7 + # Same equilibrium as the native solver. + assert abs(result.energy - w_star) < 1e-8 + assert np.linalg.norm(x - x_star) < 1e-5 + + +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_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: