vmecpp.autodiff module¶
JAX access to an in-memory VMEC++ solve and its implicit VJP.
The solver is deliberately kept outside the JAX trace. A forward call runs VMEC++ through the C++ model, while the reverse callback reruns the same model and solves the transposed interior force system. This is the usual implicit layer for a differentiable code: JAX differentiates the consumer objective, and VMEC++ supplies the producer’s residual transpose.
The first public parameterization is the fixed-boundary case, with either a
prescribed iota or a prescribed toroidal current profile (ncurr). The
differentiable parameter is one dense array with rows rbc and zbs and
shape (2, mpol, 2 * ntor + 1). This first solver wrapper deliberately
supports the stellarator-symmetric fixed-boundary case. The geometry API
itself already supports asymmetric snapshots; profile and free-boundary
parameter VJPs remain explicit unsupported cases until their residual
dependence is exposed by the exact C++ derivative path.
- class vmecpp.autodiff.DifferentiableVmec(vmec_input)¶
Bases:
objectA callable JAX view of one fixed-boundary VMEC++ input.
The current exact VJP covers the boundary coefficients. Pressure, iota, current, and flux parameters are intentionally not accepted as hidden constants: exposing them requires their residual derivatives in the C++ contract, rather than a finite-difference fallback.
- Parameters:
vmec_input (
Any)
- vmecpp.autodiff.make_solver(vmec_input)¶
Return a JAX-compatible callable that runs VMEC++ for vmec_input.
- Return type:
Example
from vmecpp import autodiff, simsopt_compat
solver = autodiff.make_solver(input) objective = lambda boundary: simsopt_compat.quasisymmetry_total(
solver(boundary), [0.6]
) value, gradient = jax.value_and_grad(objective)(boundary)
Forward execution and the VJP both invoke VMEC++ in memory. The VJP is available only in an Enzyme-enabled build, because it requires the exact transpose of the force residual. No finite-difference derivative is used.