Exact autodiff Hessian-vector product for the VMEC force

.github/actions/run-benchmarks/action.yml

@@ -31,7 +31,7 @@ runs:
       shell: bash
       run: |
         sudo apt-get update
-        sudo apt-get install -y build-essential cmake libnetcdf-dev liblapack-dev libomp-dev libhdf5-dev
+        sudo apt-get install -y build-essential cmake libnetcdf-dev libomp-dev libhdf5-dev
 
     - name: Mount bazel cache
       uses: actions/cache@v4

.github/constraints.txt

@@ -0,0 +1,3 @@
+# Block simsopt 1.11.0 wheels with unsupported AVX-512 instructions. Python 3.14
+# has no wheel and must use the corrected 1.11.0 source distribution.
+simsopt!=1.11.0; python_version != "3.14"

.github/workflows/benchmarks.yaml

@@ -8,6 +8,11 @@ permissions:
   contents: read
   pull-requests: write
 
+# CI-only dependency workaround, see .github/constraints.txt
+env:
+  PIP_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt
+  UV_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt
+
 jobs:
   benchmark:
     runs-on: ubuntu-22.04

.github/workflows/clang_tidy.yaml

@@ -28,7 +28,7 @@ jobs:
         id: review
         with:
           config_file: ".clang-tidy"
-          apt_packages: libhdf5-dev, liblapack-dev, libnetcdf-dev, gfortran, python3-dev, libomp-dev
+          apt_packages: zlib1g-dev, gfortran, python3-dev, libomp-dev
           build_dir: build
           cmake_command: cmake -B build && cmake -DCMAKE_EXPORT_COMPILE_COMMANDS:STRING=ON build
           # bazel-built test and benchmark files have no cmake compile_commands

.github/workflows/copilot-setup-steps.yml

@@ -11,6 +11,11 @@ on:
         paths:
             - .github/workflows/copilot-setup-steps.yml
 
+# CI-only dependency workaround, see .github/constraints.txt
+env:
+  PIP_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt
+  UV_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt
+
 jobs:
     # The job MUST be called `copilot-setup-steps` or it will not be picked up by Copilot.
     copilot-setup-steps:
@@ -37,10 +42,8 @@ jobs:
                   sudo apt-get install -y \
                       build-essential \
                       cmake \
-                      libnetcdf-dev \
-                      liblapack-dev \
+                      zlib1g-dev \
                       libomp-dev \
-                      libhdf5-dev \
                       python3-dev
 
             - name: Install vmecpp with test dependencies

.github/workflows/docs.yaml

@@ -16,6 +16,11 @@ concurrency:
   group: "pages"
   cancel-in-progress: false
 
+# CI-only dependency workaround, see .github/constraints.txt
+env:
+  PIP_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt
+  UV_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt
+
 jobs:
   benchmark:
     runs-on: ubuntu-22.04
@@ -48,7 +53,7 @@ jobs:
       - name: Build docs
         run: |
           sudo apt-get update
-          sudo apt-get install -y build-essential cmake libnetcdf-dev liblapack-dev libomp-dev libhdf5-dev
+          sudo apt-get install -y build-essential cmake zlib1g-dev libomp-dev
           # must actually install the package otherwise sphinx can't autogenerate the API reference
           python -m pip install -v .[docs]
           sphinx-apidoc --module-first --no-toc --force --separate -o docs/api src/vmecpp/ && sphinx-build docs html_docs

.github/workflows/full_validation.yaml

@@ -20,6 +20,11 @@ concurrency:
   group: ${{ github.workflow }}-${{ github.ref || github.run_id }}
   cancel-in-progress: ${{ github.ref_name != 'main' }}
 
+# CI-only dependency workaround, see .github/constraints.txt
+env:
+  PIP_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt
+  UV_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt
+
 jobs:
   validation:
     name: Run VMEC++ validation
@@ -28,6 +33,9 @@ jobs:
     # wout with Fortran VMEC2000 in Docker plus one with VMEC++: budget generously.
     # The short run finishes in minutes and won't come close to this limit.
     timeout-minutes: 360
+    env:
+      PIP_CONSTRAINT: ${{ github.workspace }}/vmecpp/.github/constraints.txt
+      UV_CONSTRAINT: ${{ github.workspace }}/vmecpp/.github/constraints.txt
     steps:
       - name: Check out VMEC++
         uses: actions/checkout@v4
@@ -49,8 +57,8 @@ jobs:
       - name: Install required system packages for Ubuntu
         run: |
           sudo apt-get update && sudo apt-get install -y \
-            build-essential cmake libnetcdf-dev liblapack-dev liblapacke-dev libopenmpi-dev \
-            libomp-dev libeigen3-dev nlohmann-json3-dev libhdf5-dev
+            build-essential cmake zlib1g-dev libopenmpi-dev \
+            libomp-dev libeigen3-dev nlohmann-json3-dev
 
       - name: Install vmecpp-validation's Python requirements
         run: |
@@ -61,7 +69,7 @@ jobs:
           python -m pip install -r requirements.no-vmecpp.txt
 
       - name: Install VMEC++ from this commit
-        run: python -m pip install ./vmecpp
+        run: python -m pip install "./vmecpp" "numpy<2.1"
 
       - name: Run short validation
         if: github.event_name == 'push'
@@ -96,8 +104,8 @@ jobs:
       - name: Install required packages for Ubuntu
         run: |
           sudo apt-get update && sudo apt-get install -y \
-            build-essential cmake libnetcdf-dev liblapack-dev liblapacke-dev libopenmpi-dev \
-            libomp-dev libeigen3-dev nlohmann-json3-dev libhdf5-dev
+            build-essential cmake zlib1g-dev libopenmpi-dev \
+            libomp-dev libeigen3-dev nlohmann-json3-dev
       - name: Install package
         run: python -m pip install -v .[test]
       - name: Run slow QUASR free-boundary tests

.github/workflows/pypi_publish.yml

@@ -4,18 +4,24 @@ on:
   # allows running workflows manually
   workflow_dispatch:
 
+  pull_request:
+
   release:
     types: [published]
 
+# CI-only dependency workaround, see .github/constraints.txt
+env:
+  PIP_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt
+  UV_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt
+
 jobs:
   main:
     name: Build wheels on ${{ matrix.os }}
     runs-on: ${{ matrix.os }}
     strategy:
       fail-fast: false
       matrix:
-        # ubuntu-24.04-arm not supported due to lack of netcdf-devel hdf5-devel
-        os: [ubuntu-24.04, macos-14, macos-15-intel]
+        os: [ubuntu-24.04, ubuntu-24.04-arm, macos-14, macos-15-intel]
 
     steps:
       - name: Checkout repository
@@ -26,11 +32,15 @@ jobs:
         uses: pypa/cibuildwheel@v3.3.1
         # All the settings for cibuildwheel are in pyproject.toml
         env:
+          # Linux builds run in a container, where PIP_CONSTRAINT above does not
+          # apply and the repo is mounted at /project.
+          CIBW_ENVIRONMENT_LINUX: >
+            PIP_CONSTRAINT=/project/.github/constraints.txt
+            UV_CONSTRAINT=/project/.github/constraints.txt
           CIBW_ENVIRONMENT_MACOS: >
             MACOSX_DEPLOYMENT_TARGET=${{ matrix.os == 'macos-14' && '14.0' || '15.0' }}
             DYLD_LIBRARY_PATH=/usr/local/opt/gcc/lib/gcc/current/:$DYLD_LIBRARY_PATH
             FC=gfortran-14
-            HDF5_ROOT=$(brew --prefix hdf5)
             OpenMP_ROOT=$(brew --prefix)/opt/libomp
         with:
           package-dir: .
@@ -76,7 +86,15 @@ jobs:
             # "debian:testing-20250224", # TODO(jurasic): re-enable when 3.14 support unblocked
             "fedora:41",
           ]
-    runs-on: ubuntu-24.04
+        runner: [ubuntu-24.04]
+        include:
+          # Exercise the aarch64 wheels too, on a native arm64 runner. One
+          # distro is enough here -- this job is about testing the wheel
+          # actually works on a plain distro, and the multi-distro matrix
+          # above already covers distro-specific quirks on x86_64.
+          - distro: "ubuntu:24.04"
+            runner: ubuntu-24.04-arm
+    runs-on: ${{ matrix.runner }}
     container: ${{ matrix.distro }}
 
     steps:
@@ -116,19 +134,17 @@ jobs:
           merge-multiple: true
           path: ./wheelhouse/
       - name: Install wheel
-        # Install the appropriate file for vmecpp from the wheelhouse directory.
-        # --no-index prevents installing from PyPi, so we don't accidentally install
-        # a published version instead of the newly built ones. At the same time it prevents
-        # loading dependencies from PyPi, so the first installation fails. We only use the console output
-        # from this command to get the appropriate .whl file for this platform and environment.
         run: |
           ls wheelhouse/
 
           $PYTHON -m venv venv
           . venv/bin/activate
-          wheel_filename=$(find wheelhouse/$($PYTHON --version | awk '{print $2}' | awk -F. '{printf "vmecpp*-cp%s%s-*manylinux*.whl", $1, $2}'))
+          # Let pip's own tag matching pick the single wheel compatible with this
+          # interpreter and architecture; --no-index keeps it from reaching PyPi.
+          pip download --no-index --no-deps --find-links=wheelhouse/ --dest selected-wheel/ vmecpp
+          wheel_filename=$(ls selected-wheel/*.whl)
           echo "Found matching wheel: $wheel_filename"
-          pip install $wheel_filename
+          pip install "$wheel_filename"
 
           pip install pytest
           python -m pytest tests/

.github/workflows/pyright.yaml

@@ -24,7 +24,7 @@ jobs:
       - name: Install dependencies
         run: |
           sudo apt-get update
-          sudo apt-get install build-essential cmake libnetcdf-dev liblapack-dev python3-dev
+          sudo apt-get install build-essential cmake zlib1g-dev python3-dev
       # Parse the dependencies from the pyproject.toml file and install them,
       # without installing and building the entire vmecpp project
       - run: |

.github/workflows/test_bazel.yaml

@@ -25,7 +25,7 @@ jobs:
           lfs: true  # for the mgrid test data
       - name: Install required system packages for Ubuntu
         run: |
-          sudo apt-get update && sudo apt-get install -y build-essential cmake liblapack-dev libomp-dev python3-dev python3-pip python-is-python3
+          sudo apt-get update && sudo apt-get install -y build-essential cmake libomp-dev python3-dev python3-pip python-is-python3
       - name: Mount bazel cache
         uses: actions/cache@v4
         with:

.github/workflows/test_conda_env.yaml

@@ -12,6 +12,11 @@ concurrency:
   group: ${{ github.workflow }}-${{ github.ref || github.run_id }}
   cancel-in-progress: ${{ github.ref_name != 'main' }}
 
+# CI-only dependency workaround, see .github/constraints.txt
+env:
+  PIP_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt
+  UV_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt
+
 jobs:
   test-conda-env:
     name: Test conda env

.github/workflows/test_enzyme.yaml

@@ -32,7 +32,8 @@ jobs:
           sudo ./llvm.sh "${LLVM_VERSION}"
           sudo apt-get install -y \
             "llvm-${LLVM_VERSION}-dev" "clang-${LLVM_VERSION}" "libclang-${LLVM_VERSION}-dev" \
-            cmake ninja-build
+            "libomp-${LLVM_VERSION}-dev" \
+            cmake gfortran libhdf5-dev liblapack-dev libnetcdf-dev ninja-build python3-dev
 
       - name: Cache the ClangEnzyme plugin build
         id: enzyme-cache
@@ -75,3 +76,28 @@ jobs:
             --copt="-fplugin=${ENZYME_PLUGIN}" \
             --test_output=all \
             //vmecpp/common/enzyme:enzyme_smoke_test
+
+      - uses: astral-sh/setup-uv@v6
+
+      - name: Run Enzyme force-derivative integration tests
+        env:
+          CC: clang-${{ env.LLVM_VERSION }}
+          CXX: clang++-${{ env.LLVM_VERSION }}
+          CMAKE_ARGS: >-
+            -DVMECPP_ENABLE_ENZYME=ON
+            -DVMECPP_ENZYME_PLUGIN=${{ env.ENZYME_PLUGIN }}
+        run: |
+          uv sync --reinstall-package vmecpp
+          uv run --with pytest pytest -q \
+            tests/test_hessian.py \
+            tests/test_external_optimizers.py
+
+      - name: Run exact-HVP benchmark
+        run: |
+          uv run python benchmarks/benchmark_exact_hvp.py | tee exact_hvp_benchmark.txt
+          {
+            echo "### Exact-HVP benchmark"
+            echo '```text'
+            cat exact_hvp_benchmark.txt
+            echo '```'
+          } >> "${GITHUB_STEP_SUMMARY}"

.github/workflows/tests.yaml

@@ -12,6 +12,11 @@ concurrency:
   group: ${{ github.workflow }}-${{ github.ref || github.run_id }}
   cancel-in-progress: ${{ github.ref_name != 'main' }}
 
+# CI-only dependency workaround, see .github/constraints.txt
+env:
+  PIP_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt
+  UV_CONSTRAINT: ${{ github.workspace }}/.github/constraints.txt
+
 jobs:
   pre-commit:
     env:
@@ -32,7 +37,7 @@ jobs:
         matrix:
             os: [ubuntu-22.04, macos-14]
             editable: [true, false]
-            python-version: ["3.10", "3.11", "3.12", "3.13"]
+            python-version: ["3.10", "3.11", "3.12", "3.13", "3.14"]
     steps:
       - uses: actions/checkout@v4
         with:
@@ -44,12 +49,14 @@ jobs:
       - name: Install required packages for MacOS
         if: ${{ contains(matrix.os, 'macos') }}
         run: |
-          brew install ninja libomp netcdf-cxx eigen lapack git
+          brew install ninja libomp eigen git
       - name: Install required packages for Ubuntu
         if: ${{ contains(matrix.os, 'ubuntu') }}
         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
+          # install VMEC++ deps as well as VMEC2000 deps (we need to import VMEC2000 in a test).
+          # VMEC2000 itself gets libopenblas-dev/libscalapack-mpi-dev below, which cover its
+          # LAPACK needs; VMEC++ no longer links LAPACK at all.
+          sudo apt-get update && sudo apt-get install -y build-essential cmake zlib1g-dev libomp-dev
       - name: Cache the VMEC2000 wheel
         if: ${{ matrix.os == 'ubuntu-22.04' && matrix.python-version == '3.10' }}
         uses: actions/cache@v4
@@ -111,7 +118,7 @@ jobs:
           cache: 'pip'
       - name: Install required packages for Ubuntu
         run: |
-          sudo apt-get update && sudo apt-get install -y build-essential cmake libnetcdf-dev liblapack-dev libomp-dev
+          sudo apt-get update && sudo apt-get install -y build-essential cmake zlib1g-dev libomp-dev
       - name: Install package
         run: python -m pip install .[test]
       - name: Test example scripts