examples: adjoint boundary sensitivities + SIMSOPT analytic gradient#581
Open
krystophny wants to merge 13 commits into
Open
examples: adjoint boundary sensitivities + SIMSOPT analytic gradient#581krystophny wants to merge 13 commits into
krystophny wants to merge 13 commits into
Conversation
Add a precondition flag to VmecModel.evaluate (default true, unchanged behaviour). With precondition=false the forward model 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 consistent state/gradient pair an external optimizer needs to minimise in VMEC's own basis. The native solver's preconditioned search direction (precondition=true) is a different vector; the raw gradient is the equilibrium residual and vanishes at convergence. Tests: raw force is finite and differs in direction from the preconditioned force, and drops by >1e6 from the initial guess to the converged equilibrium.
Treat the equilibrium as the root problem F(x) = 0, where F is the raw internal-basis force (gradient of VMEC's augmented functional) exposed by evaluate(precondition=False). Wire it to two solvers that reuse VMEC++'s forward model: native-style preconditioned descent and Jacobian-free Newton-Krylov (matrix-free Hessian information). Both reach the native solver's equilibrium. This is the external-differentiability path: VMEC++ as a differentiable equilibrium component an outside optimizer can drive. Quasi-Newton root-finders without a preconditioner diverge on this stiff system, which motivates exposing VMEC's preconditioner as an operator next. Tests assert both solvers reach force balance and recover the native energy and state.
Add VmecModel.apply_preconditioner(v): applies VMEC's preconditioner M^-1 (m=1, radial, lambda steps) to a vector in the decomposed basis. M^-1 is VMEC's hand-built approximate inverse Hessian; this exposes it as a reusable linear operator for preconditioned Krylov / quasi-Newton and for the Hessian solve in adjoint sensitivities. It requires a prior evaluate(precondition=true), which assembles the radial preconditioner. Validated exactly: apply_preconditioner(raw force) equals the native preconditioned search direction; the operator is linear and, once assembled, state-invariant. Use it as the inner Krylov preconditioner in Newton-Krylov: on solovev (ns=11) this cuts force evaluations from 2242 to 505 (4.4x) versus unpreconditioned JFNK, converging to the same equilibrium.
Add VmecModel.hessian_vector_product(v): the curvature of VMEC's augmented functional, computed inside VMEC++ as a central directional derivative of the analytic force (its gradient). The force is exact; only the directional step is finite-differenced. Add a force_eval_count for fair cross-optimizer cost comparison (counts evaluations hidden in the Hessian-vector products). Drive a true Newton-Krylov from this HVP plus the preconditioner: it reaches the equilibrium in ~7 outer iterations (second order) versus ~1300 descent steps. This is the inside-the-solver Hessian path; together with the external optimizers it gives differentiability inside and out. Benchmark (solovev, ns=11, force evals counted in VMEC++): preconditioned descent 2606 evals 1302 iters Newton-Krylov (JFNK) 2243 evals Newton-Krylov (preconditioned) 507 evals Newton (VMEC++ HVP + M^-1) 9194 evals 7 iters The HVP-Newton's higher force-eval count (two evals per finite-difference HVP) is what the exact Enzyme Hessian will remove.
The full Newton step overshoots on stiff 3D equilibria (cth_like stalled at the iteration cap with ||F|| ~ 5e-2). Add a backtracking line search on ||F|| so each step is damped to a decrease. With it the HVP-Newton converges on cth_like in 9 outer iterations (||F|| = 1.8e-10) and still converges solovev in 8.
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).
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.
…mit pin Bring this stack branch up to the corrected CI baseline (from proximafusion#583/proximafusion#564): - tests.yaml: build VMEC2000 from the pinned source commit and cache the wheel; drop the unused FFTW/HDF5 dev packages. - benchmarks.yaml: skip the result upload on fork PRs (read-only token). - test_simsopt_compat.py: skip vmecpp-only INDATA fields. - CMakeLists: pin abseil to the 20260107.1 commit hash for Clang >= 21.
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
1 participant
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
Use VMEC++ as a differentiable component in an external optimizer: the
boundary-shape gradient
dJ/dx_Bby the implicit-function adjoint instead offinite-differencing over boundary DOFs.
examples/vmecpp_adjoint.py: partition the state into interior/boundary,converge the interior to force balance (
solve_interior), then one adjointsolve
H_II lambda = dJ/dx_I(GMRES preconditioned byM^-1) gives the fullboundary gradient. The interior Hessian is symmetric indefinite (the lambda
constraint is a saddle), so GMRES is used, not CG.
examples/simsopt_vmec_gradient.py: a SIMSOPTOptimizablewrapping it.This PR uses the finite-difference HVP. With the exact autodiff HVP (#23) the
same adjoint gets cheaper still (numbers below).
Verification (force evals counted in VMEC++, ns=11)
The adjoint computes the same gradient as finite-differencing over the boundary
but at a cost independent of the number of boundary DOFs: 7x fewer evals on
solovev (18 DOFs), 93x on cth_like (150 DOFs) with the FD HVP, and 25x / 263x
with the exact HVP (#23). Gradients agree with the FD reference to 3.9e-4 /
4.0e-2.
Stacked on #10 (HVP).
Note (scope + follow-up): the adjoint here is exercised on the MHD energy,
whose interior cotangent vanishes at equilibrium, so the symmetric-Hessian solve
is accurate for it. VMEC's force is a scaled gradient, so
H = dF/dxisnon-symmetric; a general objective (e.g. quasisymmetry) needs the transpose
H^T. #582 adds the exact transposed Hessian-vector product and the correctedboundary gradient (forward sensitivities + an O(1) reverse adjoint), validated to
machine precision.
Tracking: #591