Exact autodiff Hessian-vector product for the VMEC force - #23
Draft
krystophny wants to merge 154 commits into
Draft
Exact autodiff Hessian-vector product for the VMEC force#23krystophny wants to merge 154 commits into
krystophny wants to merge 154 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.
The checked-in wdot histories in wout_cma.nc and wout_cth_like_free_bdy.nc were produced by the aliased in-place Eigen assignment in ComputeWOutFileContents and are not energy decay rates: they alternate small negatives with values pinned near 1.0. Patch wdot in place in both files, leaving every other variable bit-identical. The corrected values were verified to belong to the same energy history the goldens already encoded: inverting the known corruption pattern and re-applying the correct decay formula reproduces them to 2.8e-12 (cma) and 1.3e-11 (cth_like_free_bdy) max absolute difference. The code-side fix is in proximafusion#699.
* Repair shared-core consumer builds * Make wdot decay calculation alias-safe * Use the native loader path on macOS * Let Python 3.14 use the fixed SIMSOPT source release * Update test_init.py * Update test_init.py --------- Co-authored-by: Philipp Jurašić <166746189+jurasic-pf@users.noreply.github.com>
Fix dapper_type leaking into C++ bindings and non-dapper file formats
When dapper is installed, BaseModelWithNumpy aliases to dapper.DapperData,
which adds a dapper_type field to every model for type dispatch on
(de)serialization. The various places that bridge these pydantic models to
the pybind11 C++ bindings (VmecINDATA, WOutFileContents, MakegridParameters,
Mercier, JxBOutFileContents, Threed1*) iterate model_fields and getattr/
setattr against the C++ object, which has no dapper_type attribute, causing
an AttributeError. The wout NetCDF writer had the same problem via
model_dump(), which would otherwise write a spurious dapper_type variable
into .nc output.
Add a single _DAPPER_TYPE_FIELD constant plus an own_model_fields() helper
in _pydantic_numpy.py that excludes it, and route every C++-interop /
NetCDF call site through it (or through model_dump(..., exclude=...) for
the one model_dump call site). This is scoped to the literal 'dapper_type'
name only, not a generic exclusion of whatever fields BaseModelWithNumpy
happens to contribute, so it can't silently swallow unrelated fields. It's
a no-op when dapper isn't installed (BaseModelWithNumpy has no extra
fields in that case).
Also fixes several pre-existing jaxtyping shape-annotation bugs in
VmecWOut/JxBOut/Threed1* that were previously invisible: DapperData
validates each model inside a single jt.jaxtyped("context") scope, which
cross-checks that every field sharing a named dimension actually agrees
on its size. vmecpp's own fallback validator never did this cross-check:
- JxBOut.itheta/izeta were annotated 'num_half nZnT' but are allocated as
num_full in C++ (output_quantities.cc).
- Threed1GeometricAndMagneticQuantities.ygeo/yinden/yellip/ytrian/yshift
were annotated 'num_full' but are actually two num_full-length blocks
concatenated (one per toroidal symmetry plane, matching classic
Fortran VMEC's threed1 output) -- renamed to 'two_num_full'.
- Threed1AxisGeometry.raxis_asym/zaxis_asym are only populated when
lasym=True; C++ leaves them as empty vectors otherwise. Made them
Optional (None when lasym=False), matching the existing pattern for
VmecWOut's other lasym-conditional fields.
- VmecWOut.am/ac/ai share one axis ('_preset') and so must all agree on
length, but had no padding mechanism (unlike the aux profile arrays,
which already right-pad via AuxFType/AuxSType). Added the same
right-pad-to-a-common-minimum-length treatment via a new
ProfileCoeffType.
) * Add equilibrium rescaling support via HotRestartState * Fix namespace in Rescale method * Fix clang-format warnings from pre-commit * Move rescaling logic to pure Python as requested by maintainers * Auto-format python code with ruff * Fix circular import in rescale * Fix niter_array strict positivity validation by using 1 step with large ftol * Address review: optimize hot restart array lengths, use shallow copy, tighten tolerances with atol, and remove redundant comments * Scale absolute tolerances mathematically to ensure robust testing for arbitrary scale factors * Format code with ruff to fix pre-commit checks --------- Co-authored-by: dush <gandushduushguu@mail.com> Co-authored-by: Philipp Jurašić <166746189+jurasic-pf@users.noreply.github.com>
* Drop duplicated fixed-boundary JSON inputs from large test data The six fixed-boundary JSON inputs in //vmecpp_large_cpp_tests/test_data were byte-identical to the copies in //vmecpp/test_data. They were also listed in no filegroup here, so nothing could reach them: the tests in this tree already take their inputs from the "//vmecpp/test_data:..." data deps and only take the educational_VMEC reference dumps from the local test_data. regenerate_test_data.sh re-creates these files by running indata2json on the local input.* files, so they are gitignored rather than merely deleted, to keep them from coming back on the next regeneration. Deliberately untouched: - The wout_*.nc files. Despite the shared names, no two are identical: the ones here come from Fortran educational_VMEC (currumnc, currvmnc, input_extension, 1001-element aux arrays) while those in //vmecpp/test_data are VMEC++'s own output (lmns_full, fsqt, force_residual_*). Sharing them would make these tests compare VMEC++ against itself instead of against an independent reference. - The local input.* files, which carry the dump_* flags that make educational_VMEC emit the reference data. - The free-boundary files (coils.*, mgrid_solovev.nc and the solovev_free_bdy set), left as-is by request. * Update .gitignore * Drop duplicated mgrid_solovev.nc from large test data //vmecpp/test_data/mgrid_solovev.nc and the copy under //vmecpp_large_cpp_tests/test_data were byte-identical. Both were already LFS-tracked and both pointers referenced the same OID (f17260704b9f6ca96c0f7202cc99dbfe8bdb625d1103d4597fd8ca5197b96fb1), so git and LFS were already storing the 25 MB object only once; the cost was a second 25 MB copy in every checkout. The surviving copy is the one that is actually wired up: it is reached via the //vmecpp/test_data:solovev_free_bdy and :solovev_free_bdy_lforbal filegroups by output_quantities_test and vmec_in_memory_mgrid_test, and //vmecpp/test_data/solovev_free_bdy.json names that path in mgrid_file. Nothing referenced the large-tree copy: it sat in no filegroup, because the solovev_free_bdy case was never wired up in this tree. The unused solovev_free_bdy.json here still names the removed path, but it is equally unwired, and regenerate_test_data.sh rebuilds the mgrid from coils.solovev via makegrid if that case is ever enabled.
…afusion#600) * vmec_indata: warn when the input boundary is spectrally dense Fixed-boundary runs whose boundary carries a lot of high-poloidal-mode content tend to fail to initialize or converge poorly, and the only signal today is a late, generic solver abort. IsConsistent now computes the boundary spectral width <M> (the wout.specw / ComputeSpectralWidth quantity, evaluated on the input coefficients) and warns above a calibrated threshold before the solve, for fixed-boundary runs only. The threshold (6) is set from ripple sweeps on a tokamak and a 3D stellarator: fixed-boundary runs converge for <M> below ~6 and fail from ~7.2, while all bundled inputs stay below 1.75. * Update src/vmecpp/cpp/vmecpp/common/vmec_indata/vmec_indata.cc Co-authored-by: Philipp Jurašić <166746189+jurasic-pf@users.noreply.github.com> * Share one spectral-width implementation between the geometry and the boundary check Extract the per-surface spectral width out of FourierGeometry::ComputeSpectralWidth into common/spectral_width and call it from both the geometry and a new Boundaries::ComputeSpectralWidth. The boundary warning moves to Boundaries::setupFromIndata, which is the lowest layer that can reach the shared implementation, and its threshold is 2.0. * Keep the shared spectral width in the Fourier basis target instead of its own The function weights coefficients by the mscale and nscale that FourierBasisFastPoloidal defines, so it belongs beside them; both callers already depended on that target. Drops the separate common/spectral_width library and its CMake entry. --------- Co-authored-by: Philipp Jurašić <166746189+jurasic-pf@users.noreply.github.com>
…fusion#702) VmecWOut.from_wout_file() builds a dict keyed by the wout file's native variable names (e.g. lasym__logical__) and validates it with model_validate(attrs, by_alias=True), expecting the by_alias override to apply. But VmecWOut.model_config set validate_by_alias=False, and with a model_validator(mode="wrap") present (used to scope jaxtyping axis checks to the model), a runtime by_alias/by_name override passed to model_validate() is silently dropped -- a known pydantic-core bug (pydantic/pydantic#13661), fixed upstream but not yet in a pinned release. Set validate_by_alias=True directly in VmecWOut.model_config instead of relying on the call-time override, which works correctly regardless.
…::PartialPivLU (proximafusion#706) vmecpp's calls exactly two LAPACK routines, dgetrf_ and dgetrs_ in LaplaceSolver. Because of that, we bundle a full LAPACK library with the wheel. This PR replaces both calls with Eigen::PartialPivLUEigen::MatrixXd, removing one more dependency. Benchmarks continue to report similar numbers.
…sion#704) ## Problem ~76% of the `vmecpp` wheel is transitive shared libs. The biggest chunk is netCDF's OPeNDAP logic, for remote access which we don't even use. Reduces wheel size 13MB -> 9MB ## Root cause `[tool.cibuildwheel.linux].before-build` installed manylinux_2_28's distro `netcdf-devel`/`hdf5-devel` packages. Those distro packages are themselves built *with* DAP/OPeNDAP support enabled, which is what drags in libcurl and its whole TLS/Kerberos/LDAP dependency closure. The repo's Bazel build (`src/vmecpp/cpp/third_party/hdf5/BUILD.bazel`, `third_party/netcdf4/BUILD.bazel`) already builds both libraries from source, static, with DAP fully disabled -- this is a proven, working recipe already used for local dev/test builds - but it was never wired into the CMake-based wheel-build path.
Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: jurasic-pf <166746189+jurasic-pf@users.noreply.github.com>
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
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.