Skip to content

pybind: expose the unpreconditioned internal-basis gradient#577

Open
krystophny wants to merge 7 commits into
proximafusion:mainfrom
itpplasma:expose-internal-gradient
Open

pybind: expose the unpreconditioned internal-basis gradient#577
krystophny wants to merge 7 commits into
proximafusion:mainfrom
itpplasma:expose-internal-gradient

Conversation

@krystophny

@krystophny krystophny commented Jun 14, 2026

Copy link
Copy Markdown
Contributor

Stacked PR — part 14/19 of the differentiable-VMEC++ series. based directly on main.
Diff is cumulative (includes ancestor commits) because the branches are stacked on the fork; review the net change described below.


What

Add a precondition flag to VmecModel.evaluate (default true, so existing
behaviour is unchanged). With precondition=false the forward model returns at
the existing INVARIANT_RESIDUALS checkpoint, so get_forces() returns the raw,
unpreconditioned force.

Why

This is the first step toward driving VMEC++ with gradient- and Hessian-based
optimizers (inside the solver and as a differentiable component of an external
stellarator optimizer). Such optimizers need a consistent (state, gradient) pair.

VMEC minimizes an augmented functional: the MHD energy plus the Hirshman
spectral-condensation constraint and the lambda (magnetic-coordinate) constraint.
Its "force" is the gradient of that augmented functional with respect to the
decomposed, internal-basis state, where the scalxc scaling makes the state and
force conjugate. The raw force at the INVARIANT_RESIDUALS checkpoint is exactly
that gradient.

The native solver instead exposes the preconditioned search direction
(precondition=true), which is a different vector: the preconditioner is a
non-trivial metric, not a scalar. An external optimizer minimizing in VMEC's
basis needs the raw gradient, plus VMEC's preconditioner applied separately as
the metric (a later patch).

No core solver code changes; only the pybind layer selects the checkpoint.

Verification

evaluate(1, 2) is the documented cold-start case (forces initialised to 1.0);
the raw-gradient path uses iter1 >= 2, where the force is well defined.

Empirically, on examples/data/solovev.json (ns=11), the raw force is the
equilibrium residual: it vanishes once the native solver converges, and it is
not the preconditioned direction.

initial guess : ||raw force|| = 1.718e-01
converged     : ||raw force|| = 5.093e-09     (3.4e7 reduction)
cos(raw, preconditioned) = -0.10             (different direction, not scale)

New tests (tests/test_internal_gradient.py):

test_raw_force_differs_from_preconditioned PASSED
test_raw_force_vanishes_at_equilibrium     PASSED
test_cold_start_is_excluded                PASSED
3 passed

Tracking: #590

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.
Satisfies the docformatter pre-commit hook (was failing CI).
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.
@krystophny krystophny marked this pull request as draft June 15, 2026 04:48
…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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant