Exact autodiff Hessian-vector product for the VMEC force#23
Draft
krystophny wants to merge 103 commits into
Draft
Exact autodiff Hessian-vector product for the VMEC force#23krystophny wants to merge 103 commits into
krystophny wants to merge 103 commits into
Conversation
Add the implicit-function adjoint that turns VMEC++ into a gradient-providing equilibrium component for SIMSOPT, the original goal. vmecpp_adjoint.py: for a converged fixed-boundary equilibrium F_I(x)=0, the boundary sensitivity of a scalar objective J follows from H_II lambda = dJ/dx_I, dJ/dx_B = dJ/dx_B - (dF_I/dx_B)^T lambda, with H the symmetric Hessian of the augmented functional. It is matrix-free via hessian_vector_product and apply_preconditioner (the SPD interior system is solved with preconditioned CG). One Hessian solve gives the whole boundary gradient, versus one equilibrium re-solve per boundary DOF for finite differences. simsopt_vmec_gradient.py: VmecEnergy wraps this as a SIMSOPT Optimizable whose dJ is the adjoint gradient, plus a gradient-cost benchmark. Verified: the adjoint gradient matches brute-force re-solve finite differences (rel 2.4e-4) and the SIMSOPT Optimizable's dJ matches finite differences of J (rel ~1e-6). On solovev (ns=11, 18 boundary DOFs) the adjoint boundary gradient costs 762 force evaluations versus 9112 for finite differences (12x), and the gap grows with the boundary DOF count.
Two correctness fixes for stiff 3D equilibria (cth_like): - VMEC's augmented-Lagrangian Hessian is symmetric *indefinite* (the lambda constraint makes it a saddle, not a minimum), so CG silently gives the wrong adjoint there. Use GMRES, which handles indefinite systems, for the H_II solve and the interior Newton solve. With a loose, restarted tolerance the adjoint solve stays cheap. - Add a backtracking line search to solve_interior so the interior re-solve (used by the SIMSOPT wrapper and the finite-difference reference) converges on 3D instead of overshooting. Verified with a directional-derivative check against a re-converged finite-difference reference: solovev 1.5e-4, cth_like 2.2e-2 relative; both previously agreed only in 2D. Boundary-gradient cost on solovev: 626 force evaluations (analytic adjoint) versus 10460 (finite differences).
Move the composed local force map g (MHD force chain + hybrid lambda force) into local_force_composition.h (ComputeLocalForceDensity), parameterized by the radial partition offsets so the same composition serves the Enzyme autodiff test and the exact Hessian-vector product over the live model state. The autodiff test now calls the shared composition; forward/reverse vs finite-diff and forward/reverse agreement are unchanged (2.55e-8, 4e-15).
Add exact_force_jvp.cc/.h: ExactForceDensityJvp wraps one Enzyme forward pass over ComputeLocalForceDensity, returning the force-density tangent for a geometry tangent. This translation unit is compiled with the Clang/Enzyme plugin and is the single nonlinear pass the exact Hessian-vector product calls; the rest of VMEC++ stays normally compiled and wraps it with the linear spectral transforms. The autodiff test now also validates this standalone entry point against a finite difference of the composition's force-density output.
Add IdealMhdModel::applyExactForceJacobian: given the packed real-space geometry primal and a geometry tangent, differentiate the MHD-plus-lambda force density by one Enzyme forward pass (ExactForceDensityJvp), scatter the tangent into the real-space force members, then apply the linear forward transform and preconditioner decomposition (forcesToFourier, decomposeInto, m1Constraint, zeroZForceForM1) exactly as the tail of update() does. The geometry tangent is supplied by the caller via the linearity of geometryFromFourier (geom(x+v) - geom(x)), so the only nonlinear step is the single Enzyme pass. The constraint force (a linear Fourier bandpass over a nonlinear product) is omitted here and added separately. Compiles under clang-21; end-to-end exact-vs-finite-difference validation of the assembled Hessian-vector product against VmecModel runs once the internal solver stack (VmecModel HVP + preconditioner) is merged with this kernel stack.
Wire the exact force Hessian-vector product through to Python: VmecModel.exact_hessian_vector_product(v) computes H v = T^T J_g T v with one Enzyme forward pass over the local force-density composition, using the linearity of geometryFromFourier for the exact geometry tangent (geom(x+v) - geom(x)) and the existing forward transform + decomposition for the output. exact_force_jvp.cc is compiled into the core library with the Enzyme plugin; VMECPP_ENABLE_ENZYME guards the callers so the default plugin-free build is unchanged. Built into the extension with clang-21 + Enzyme. The result is 96% cosine aligned with the finite-difference HVP on solovev; the remaining difference is the spectral-condensation constraint force, which is omitted from this pass and added next (it carries a linear Fourier bandpass over a nonlinear product).
Pointer-ize the constraint-force bandpass (ComputeDeAliasConstraintForce in constraint_force_kernel.h, explicit reductions so it differentiates under Enzyme) and have the free function deAliasConstraintForce call it; the solver energy stays bit-exact at 1 and 4 threads. Extend ComputeLocalForceDensity with an optional constraint stage (geometry blocks 16-19 = rCon/zCon/ruFull/zuFull, force blocks 16-19 = frcon/fzcon) and enable it in applyExactForceJacobian, with VmecModel packing the constraint geometry. Status: the exact HVP runs end-to-end and is 96% cosine-aligned with the finite-difference HVP on solovev. Including the constraint changes the result only marginally, so the remaining ~32% magnitude/direction difference is in the MHD/lambda or transform path, not the constraint; reaching exact==FD needs tracing the residual geometry dependence (e.g. the iter==iter rCon0 volume extrapolation), which is the next debugging step before benchmarking.
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| void ExactForceDensityJvp(const double* geom, const double* dgeom, double* work, | ||
| double* dwork, double* force, double* dforce, | ||
| const LocalForceComposition* c) { | ||
| __enzyme_fwddiff<void>((void*)ComputeLocalForceDensity, enzyme_dup, geom, |
There was a problem hiding this comment.
warning: C-style casts are discouraged; use reinterpret_cast [google-readability-casting]
Suggested change
| __enzyme_fwddiff<void>((void*)ComputeLocalForceDensity, enzyme_dup, geom, | |
| __enzyme_fwddiff<void>(reinterpret_cast<void*>(ComputeLocalForceDensity), enzyme_dup, geom, |
| double* gsc = s; | ||
| s += c->ntor + 1; | ||
| double* gcs = s; | ||
| s += c->ntor + 1; |
There was a problem hiding this comment.
warning: Value stored to 's' is never read [clang-analyzer-deadcode.DeadStores]
s += c->ntor + 1;
^Additional context
src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/local_force_composition.h:193: Value stored to 's' is never read
s += c->ntor + 1;
^… HVP Add composedForceResidual (exposed as VmecModel.composed_force_residual): the max difference between the flat-buffer force-density composition and the production force-density members at the current state. It is exactly 0.0 on solovev and cth_like, proving the composition (all kernels + constraint + bandpass) reproduces the production force density bit-for-bit. Gate zeroZForceForM1 in applyExactForceJacobian on fsqz < 1e-6, matching the tail of update(). Diagnosis of the end-to-end exact-vs-FD HVP gap (~27%, eps-independent, and the exact operator is linear to 0.4%): with the composition proven exact and the forward transform matching update(), the residual is isolated to the geometry-tangent / autodiff-of-the-2D-path. The microtest validates the Enzyme JVP only for the 3D composition; the 2D (axisymmetric) JVP and the linearity of the packed geometry tangent are the next things to instrument.
…==FD Two fixes make the exact Hessian-vector product match the finite-difference HVP to finite-difference precision: 1. Geometry tangent by small central difference. geom(decomposed_x) is only near-linear (the preconditioner scaling is frozen per step but the transform still has curvature), so the unit-step difference geom(x+v)-geom(x) contaminated the tangent with higher-order terms (~30% error). Use a small central step; the nonlinear force kernels are still differentiated exactly by the single Enzyme pass. 2. Compute the constraint reference rCon0/zCon0 inside the composition by the rzConIntoVolume extrapolation (rCon0[jF] = rCon[LCFS] * s_full) instead of freezing it. It is linear in the geometry and recomputed every step by the solver, so freezing it left a ~3% residual. Validation (exact vs FD HVP, eps-independent): solovev 3.9e-6 (cos 1.000000), cth_like 4.2e-3 (cos 0.999991). solovev (ncurr=0) is exact to the FD floor; the cth_like residual is the ncurr=1 chi' = f(geometry) term in computeBContra, which is still frozen and is the next term to differentiate. Adds isolation diagnostics force_density_jvp_residual and (earlier) composed_force_residual.
solve_newton_exact_hvp drives the globalized preconditioned Newton-Krylov with VmecModel.exact_hessian_vector_product (one Enzyme pass) instead of the finite-difference HVP. Requires an Enzyme-enabled build.
krystophny
changed the title
…pdate Replace the finite-difference geometry tangent (and its full update() calls) with the exact linear pre-chain applied directly to the perturbation: IdealMhdModel::packGeometry runs decomposeInto + m1Constraint + extrapolate + geometryFromFourier on a vector and packs the 20-block geometry with the computeBContra lambda normalization. Applied to the state it gives the geometry (primal, with phipF on lu_e); applied to v it gives the exact geometry tangent, since the chain is linear. The only nonlinear step is the single Enzyme pass. The exact HVP is now fully analytic/autodiff (no finite difference anywhere) and does no full force evaluation per matvec. Accuracy is unchanged (exact vs FD HVP 3.9e-6 solovev, 4.2e-3 cth_like). Force-evaluation count of the exact-HVP Newton drops from 14136 to 27 on solovev (52 on cth_like), since matvecs no longer call update(); wall-clock is still dominated by the GMRES matvecs (each a transform + Enzyme pass), where preconditioned JFNK remains faster.
The primal geometry depends only on the state, not the Krylov vector, so a GMRES solve recomputed it on every matvec. Cache it (invalidated by SetState), halving the geometry-transform work per matvec. Accuracy unchanged (exact vs FD 3.9e-6 solovev); wall-clock drops (cth_like internal 25.7->21.8s, adjoint 11.5->10.4s).
For a constrained toroidal-current profile (ncurr==1) chi' is recomputed from the geometry every step (chi' = (currH - jvPlasma)/avg_guu_gsqrt with jvPlasma, avg_guu_gsqrt surface integrals of the metric and field), so freezing it left a 0.4% residual on cth_like. Compute chi' inside the composition from the pre-chi' contravariant field; ncurr==0 keeps the frozen iota*phi' profile. Exact vs FD HVP now: solovev 3.9e-6, cth_like 6.6e-6 (was 4.2e-3) - exact to the finite-difference floor on both cases.
This was referenced Jun 14, 2026
benchmark_exact_hvp.py: preconditioned JFNK vs FD-HVP vs exact-HVP Newton-Krylov (internal solver). benchmark_adjoint_gradient.py: FD-over-boundary vs FD-HVP adjoint vs exact-HVP adjoint (external/SIMSOPT boundary gradient). These produce the performance tables in the PR descriptions; the exact-HVP rows need an Enzyme-enabled build.
boundary_gradient and _interior_operators take exact=True to use exact_hessian_vector_product (no force evaluation per matvec); the SIMSOPT VmecEnergy.gradient wrapper uses it by default. A real SIMSOPT least_squares_serial_solve loop (minimize (energy-target)^2 over the boundary) converges with this analytic gradient to |J-target|/target = 7e-7 in 10 iterations.
…ts JFNK The inner Krylov solve dominated wall-clock: the augmented Hessian is indefinite, so a fixed tight inner tolerance made GMRES take ~580 matvecs per Newton step even though only ~5-10 outer steps are needed. Profiling: 8 Newton iters but 4655 matvecs, one exact HVP matvec is 0.04 ms, and only 21% of the time was in the HVP. Fix: Eisenstat-Walker adaptive inner forcing (loose-early/tight-late) and lgmres (recycles Krylov vectors), applied to both Newton-HVP solvers (both already preconditioned by VMEC's M^-1). Final, fair comparison (all preconditioned, ns=11): solovev: precond JFNK 507 ev / 0.08 s ; exact-HVP Newton 17 ev / 0.03 s cth_like: precond JFNK 1633 ev / 2.00 s; exact-HVP Newton 26 ev / 1.32 s The exact-autodiff-HVP Newton-Krylov now beats preconditioned JFNK on both cases in both force evaluations (30x / 63x fewer) and wall-clock (2.7x / 1.5x faster). Also add examples/simsopt_optimization_loop.py: a real SIMSOPT least_squares_serial_solve loop driving the boundary to a target energy via the analytic adjoint gradient (converges to |J-target|/target = 7e-7).
…t fallback Apply pre-commit formatting across the touched files. Make the SIMSOPT VmecEnergy.gradient auto-detect the exact HVP: use exact_hessian_vector_product when the extension was built with Enzyme, otherwise fall back to the finite-difference HVP, so the default (plugin-free) build -- and test_simsopt_gradient -- pass without an Enzyme build.
…mhd_model ideal_mhd_model.cc includes the header-only force kernels and the autodiff composition/JVP header; the bazel sandbox needs them declared or the bazel builds (opt/asan/ubsan/tsan, test_bazel) fail with 'jacobian_kernel.h: No such file'. Add them via glob so each branch picks up whatever exists on it. The Enzyme JVP .cc stays CMake-only (guarded by VMECPP_ENABLE_ENZYME).
The 'Compare benchmark result' step uses github-action-benchmark with comment-on-alert and the GITHUB_TOKEN, which is read-only for pull requests from forks -> 'Resource not accessible by integration'. Gate that step on the PR coming from the same repo so fork PRs still run the benchmarks but skip the write-back instead of failing.
The pinned vmec-0.0.6 cp310 wheel was f90wrapped against numpy 1.x. Under the numpy 2.x that the test env now resolves, importing it dies in the f90wrap array interface (f90wrap_vmec_input__array__rbc: 0-th dimension must be fixed to 2 but got 4), so test_ensure_vmec2000_input_from_vmecpp_input could never actually run on CI (and is currently red on main too, where the wheel's runtime libs are not even installed). Build VMEC2000 from upstream source with current f90wrap, which produces numpy-2-compatible bindings. The recipe mirrors SIMSOPT's own CI (hiddenSymmetries/VMEC2000, cmake/machines/ubuntu.json). An explicit 'import vmec' check in the install step surfaces any remaining problem here rather than as a confusing test failure.
With VMEC2000 built from current upstream source, the compatibility test runs for the first time and hits vmecpp indata fields that have no counterpart in the legacy VMEC2000 INDATA namelist (e.g. free_boundary_method), which raised AttributeError. The test explicitly checks only the common subset, so guard the lookup with hasattr and skip fields VMEC2000 does not have, instead of enumerating them one by one.
Fix typos in docs
Run every solver once on explicit Solov'ev and CTH-like inputs. Keep the strict internal-state comparison on the reproducible 2D case, and verify force balance and energy in 3D without optional diagnostics or repeated evaluations. Co-authored-by: Philipp Jurašić <166746189+jurasic-pf@users.noreply.github.com>
Add the implicit-function adjoint that turns VMEC++ into a gradient-providing equilibrium component for SIMSOPT, the original goal. vmecpp_adjoint.py: for a converged fixed-boundary equilibrium F_I(x)=0, the boundary sensitivity of a scalar objective J follows from H_II lambda = dJ/dx_I, dJ/dx_B = dJ/dx_B - (dF_I/dx_B)^T lambda, with H the symmetric Hessian of the augmented functional. It is matrix-free via hessian_vector_product and apply_preconditioner (the SPD interior system is solved with preconditioned CG). One Hessian solve gives the whole boundary gradient, versus one equilibrium re-solve per boundary DOF for finite differences. simsopt_vmec_gradient.py: VmecEnergy wraps this as a SIMSOPT Optimizable whose dJ is the adjoint gradient, plus a gradient-cost benchmark. Verified: the adjoint gradient matches brute-force re-solve finite differences (rel 2.4e-4) and the SIMSOPT Optimizable's dJ matches finite differences of J (rel ~1e-6). On solovev (ns=11, 18 boundary DOFs) the adjoint boundary gradient costs 762 force evaluations versus 9112 for finite differences (12x), and the gap grows with the boundary DOF count.
Two correctness fixes for stiff 3D equilibria (cth_like): - VMEC's augmented-Lagrangian Hessian is symmetric *indefinite* (the lambda constraint makes it a saddle, not a minimum), so CG silently gives the wrong adjoint there. Use GMRES, which handles indefinite systems, for the H_II solve and the interior Newton solve. With a loose, restarted tolerance the adjoint solve stays cheap. - Add a backtracking line search to solve_interior so the interior re-solve (used by the SIMSOPT wrapper and the finite-difference reference) converges on 3D instead of overshooting. Verified with a directional-derivative check against a re-converged finite-difference reference: solovev 1.5e-4, cth_like 2.2e-2 relative; both previously agreed only in 2D. Boundary-gradient cost on solovev: 626 force evaluations (analytic adjoint) versus 10460 (finite differences).
Separate the legacy previous-residual m=1 projection used by native iteration from the exact constrained projection used by external force evaluations. Cover both directions across the residual threshold.
Use the low-level VMEC model for the native benchmark row so every solver reports the same force norm and MHD functional. Count SciPy Newton-Krylov outer iterations through its callback and correct the preconditioner description.
# Conflicts: # examples/external_optimizers.py # src/vmecpp/cpp/vmecpp/vmec/pybind11/pybind_vmec.cc # tests/test_external_optimizers.py # tests/test_hessian.py
…integration # Conflicts: # examples/simsopt_vmec_gradient.py # examples/vmecpp_adjoint.py # tests/test_adjoint.py # tests/test_simsopt_gradient.py
…exact-hvp-integration
…' into exact-hvp-integration # Conflicts: # examples/external_optimizers.py
…' into exact-hvp-integration # Conflicts: # src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/ideal_mhd_model.cc # src/vmecpp/cpp/vmecpp/vmec/pybind11/pybind_vmec.cc
Fix formatting in docstring for clarity.
* Make raw VMEC forces history independent Separate the legacy previous-residual m=1 projection used by native iteration from the exact constrained projection used by external force evaluations. Cover both directions across the residual threshold. * Use compact storage for constraint policy * Name the m=1 gauge policy explicitly
# Conflicts: # src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/ideal_mhd_model.cc # src/vmecpp/cpp/vmecpp/vmec/ideal_mhd_model/ideal_mhd_model.h # src/vmecpp/cpp/vmecpp/vmec/pybind11/pybind_vmec.cc
examples/data/cth_like_fixed_bdy.out.h5 and examples/data/solovev.out.h5 were run artifacts written next to the input files when the examples were executed, committed by mistake in 267fdd0. Nothing references them; the canonical test equilibria live in src/vmecpp/cpp/vmecpp_large_cpp_tests/test_data/.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
End-to-end exact, finite-difference-free Hessian-vector product for VMEC++'s
force, from the shared force-density kernels through Enzyme to Python, plus the
solvers and adjoint that use it. Integrates the kernel/autodiff stack (#13-#22)
with the internal-solver and adjoint stacks (#7-#11).
VmecModel.exact_hessian_vector_product(v)computesH v = Tᵀ J_g T v:T v: the linear pre-chain applied directly tov--exact, no finite difference, no full
update()(packGeometry); the primalgeometry is cached across Krylov matvecs.
J_g: one Enzyme forward pass over the full nonlinear force density(MHD + lambda + spectral-condensation constraint with its bandpass and
rzConIntoVolumereference, + thencurr=1chi' recomputation).Tᵀ:forcesToFourier+ preconditioner decomposition.VMECPP_ENABLE_ENZYMEguards the callers; the default build is unchanged.Verification: exact vs finite-difference HVP (eps-independent)
Exact to the FD floor on both.
composed_force_residual= 0.0 (composition ==production force density bit-for-bit); Enzyme JVP vs FD 2D and 3D (2.4e-8); plugin
entry point (8.6e-9). Solver stays bit-exact at 1 and 4 threads.
Internal solver (force evals counted in VMEC++, ns=11; all preconditioned by M^-1, Eisenstat-Walker forcing, lgmres)
The exact-autodiff-HVP Newton-Krylov beats best-of-breed preconditioned JFNK on
both cases, in both metrics: 30x / 63x fewer force evaluations and 2.7x / 1.5x
less wall-clock. Each matvec is a cheap transform + one Enzyme pass (no force
evaluation), and Eisenstat-Walker forcing keeps the indefinite inner solve to a
handful of matvecs early on.
External / SIMSOPT (boundary-shape gradient dJ/dx_B)
The exact-HVP adjoint computes the boundary gradient with 25x (solovev) / 263x
(cth_like) fewer force evaluations than finite-differencing over the boundary,
at a cost independent of the DOF count. A real SIMSOPT optimization loop
(
examples/simsopt_optimization_loop.py) drives the boundary to a target energywith this analytic gradient via
least_squares_serial_solveand converges to|J-target|/target = 7e-7.Conclusion
VMEC++ is now an exactly-differentiable equilibrium component:
case and metric measured;
than finite differences, exact to ~1e-6, in a working optimization loop.
FD status: gradient, preconditioner, and HVP are all analytic/autodiff -- zero
finite difference (the only FD left is inside scipy's JFNK, a third-party method).
Reproduce:
examples/benchmark_exact_hvp.py,examples/benchmark_adjoint_gradient.py,examples/simsopt_optimization_loop.py(Enzyme-enabled build). Base #10.