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 solve pins the m=1 poloidal-origin gauge (VmecInput.always_fix_m1_gauge): the native iteration lets that gauge drift under its force until fsqz drops below 1e-6 and freezes it wherever it is, so the converged state would depend on the iteration history and the fixed-gauge force Jacobian would not describe it. With the gauge pinned, the gauge entries of the state are a linear function of the boundary, and the VJP pulls them back to the boundary like the boundary entries themselves.

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. Power-series profile parameters are differentiable as well; free-boundary parameters are not.

class vmecpp.autodiff.DifferentiableVmec(vmec_input)

Bases: object

A callable JAX view of one fixed-boundary VMEC++ input.

The exact VJP covers the boundary coefficients and the power-series profile parameters passed as profiles (see autodiff_wout.PROFILE_PARAMETERS).

Parameters:

vmec_input (Any)

vmec_input: Any

property parameter_shape: tuple[int, int, int]

property ns: int

property geometry_size: int

property extra_size: int

Size of the non-differentiated outputs that follow the geometry.

property output_shape: tuple[int]

vmecpp.autodiff.make_solver(vmec_input)

Return a JAX-compatible callable that runs VMEC++ for vmec_input.

Return type:

DifferentiableVmec

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.