pybind: Hessian-vector product inside VMEC++ + internal Newton-Krylov#10
Draft
krystophny wants to merge 12 commits into
Draft
pybind: Hessian-vector product inside VMEC++ + internal Newton-Krylov#10krystophny wants to merge 12 commits into
krystophny wants to merge 12 commits into
Conversation
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.
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.
# Conflicts: # examples/external_optimizers.py # tests/test_external_optimizers.py # tests/test_internal_gradient.py
make_model now reguesses the magnetic axis when the initial geometry has a singular Jacobian, mirroring the native solver's first-iterate axis reguess (vmec.cc SolveEquilibriumLoop). Inputs that ship no axis (raxis/zaxis all zero, e.g. cma.json) otherwise return a zero raw force at the BAD_JACOBIAN checkpoint. Add a cma test (3D stellarator, nfp=2, ntor=6) that exercises the non-axisymmetric force chain: after the reguess the raw internal-basis force and the Hessian-vector product are finite and nonzero.
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
Expose VMEC++'s force Hessian-vector product through pybind
(
VmecModel.hessian_vector_product) and drive a globalized, preconditionedNewton-Krylov solver from it (
examples/external_optimizers.py: solve_newton_hvp). Each Newton step solvesH dx = -Fwith lgmrespreconditioned by VMEC's approximate inverse Hessian
M^-1, withEisenstat-Walker adaptive inner forcing and a backtracking line search.
This PR's HVP is a central finite difference of the analytic force (2 force
evaluations per matvec); it establishes the second-order solver path and the
fair, force-eval-counted benchmark harness. The exact, finite-difference-free
autodiff HVP that replaces it is PR #23.
Verification (force evals counted in VMEC++, ns=11)
All methods are preconditioned by
M^-1and converge to the native equilibriumenergy. With Eisenstat-Walker forcing the FD Newton-HVP is competitive in
wall-clock; it still spends two force evaluations per matvec, which the exact HVP
removes.
Conclusion
This is the second-order path with the finite-difference HVP. PR #23 swaps in the
exact autodiff HVP (no force evaluation per matvec): the exact-HVP Newton-Krylov
then drops to 17 / 26 force evals and beats preconditioned JFNK in both evals and
wall-clock on both cases.
Stacked on #9 (preconditioner).