API reference¶

The toroidal production API is vmex.core; the open-field-line API is vmex.mirror. Modules are grouped as in Architecture; every docstring names the VMEC2000 counterpart it ports.

Inputs and profiles¶

VMEC input handling: the &INDATA namelist and VMEC++-style JSON.

VMEC2000 counterparts: LIBSTELL/Sources/Modules/vmec_input.f (read_indata_namelist: variable set and defaults) and readin.f (post-read normalizations). The JSON schema follows VMEC++ (vmecpp.VmecInput): identical key names, boundary coefficients as sparse {"m": int, "n": int, "value": float} lists, dense axis arrays, and adiabatic_index accepted as an alias for gamma.

VmecInput is a frozen dataclass holding the full INDATA content this code base consumes, with VMEC2000 defaults. Parsing is host-side NumPy code (nothing here needs JAX).

Normalizations applied on construction (all from VMEC2000):

read_indata_namelist: raxis_s[0] = 0 and zaxis_s[0] = 0; the obsolete RAXIS/ZAXIS arrays override RAXIS_CC/ZAXIS_CS where nonzero; niter_array falls back to NITER when absent.
readin.f: lfreeb is forced False when mgrid_file == 'NONE'; nvacskip <= 0 falls back to nfp.
Boundary coefficients outside |n| <= ntor, 0 <= m < mpol are dropped (VMEC2000 reads them into oversized arrays but never uses them).

Index conventions: rbc/zbs/rbs/zbc are dense 2D arrays of shape (2*ntor + 1, mpol) indexed [n + ntor, m], i.e. rbc[n + ntor, m] is the INDATA coefficient RBC(n, m).

class vmex.core.input.VmecInput(lasym: bool = False, nfp: int = 1, mpol: int = 6, ntor: int = 0, ntheta: int = 0, nzeta: int = 0, ns_array: Any = None, ftol_array: Any = None, niter_array: Any = None, delt: float = 1.0, tcon0: float = 1.0, aphi: Any = None, phiedge: float = 1.0, nstep: int = 10, pmass_type: str = 'power_series', am: Any = None, am_aux_s: Any = None, am_aux_f: Any = None, pres_scale: float = 1.0, gamma: float = 0.0, spres_ped: float = 1.0, ncurr: int = 0, pcurr_type: str = 'power_series', ac: Any = None, ac_aux_s: Any = None, ac_aux_f: Any = None, curtor: float = 0.0, piota_type: str = 'power_series', ai: Any = None, ai_aux_s: Any = None, ai_aux_f: Any = None, bloat: float = 1.0, raxis_c: Any = None, zaxis_s: Any = None, raxis_s: Any = None, zaxis_c: Any = None, rbc: Any = None, zbs: Any = None, rbs: Any = None, zbc: Any = None, lfreeb: bool = True, mgrid_file: str = 'NONE', extcur: Any = None, nvacskip: int = 1, mfilter_fbdy: int = -1, nfilter_fbdy: int = -1, precon_type: str = 'NONE', prec2d_threshold: float = 1e-30)¶

Full &INDATA content with VMEC2000 semantics and defaults.

Defaults are the initializations in read_indata_namelist (vmec_input.f), after the readin.f normalizations documented in the module docstring. Array fields are NumPy arrays; None defaults are resolved in __post_init__ (they depend on mpol/ntor).

lasym: bool = False¶: non-stellarator-symmetric mode

nfp: int = 1¶: number of field periods

mpol: int = 6¶: poloidal modes m = 0..mpol-1

ntor: int = 0¶: toroidal modes n = -ntor..ntor

ntheta: int = 0¶: poloidal grid points (0 -> VMEC default)

nzeta: int = 0¶: toroidal grid points (0 -> VMEC default)

ns_array: Any = None¶: radial surfaces per stage (default [31])

ftol_array: Any = None¶: force tolerance per stage (default [1e-10])

niter_array: Any = None¶: iteration cap per stage (default [100] = NITER)

delt: float = 1.0¶: initial time step

tcon0: float = 1.0¶: constraint-force multiplier (bcovar.f)

aphi: Any = None¶: radial-flux remap polynomial (default [1,0,…], len 20)

phiedge: float = 1.0¶: total enclosed toroidal flux [Wb]

nstep: int = 10¶: iterations between progress prints

am: Any = None¶: pmass coefficients (dense, len >= 21)

am_aux_s: Any = None¶: pmass spline knots s

am_aux_f: Any = None¶: pmass spline values

pres_scale: float = 1.0¶: pressure scale factor [Pa]

gamma: float = 0.0¶

also ‘adiabatic_index’)

Type:: adiabatic index (JSON

spres_ped: float = 1.0¶: pressure pedestal s (profil1d.f clamp)

ncurr: int = 0¶

prescribed current

Type:: 0
Type:: prescribed iota, 1

ac: Any = None¶: pcurr coefficients (dense, len >= 21)

curtor: float = 0.0¶: total toroidal current [A]

ai: Any = None¶: piota coefficients (dense, len >= 21)

bloat: float = 1.0¶: profile-argument expansion factor

raxis_c: Any = None¶: R axis cos coefficients (INDATA RAXIS_CC)

zaxis_s: Any = None¶: Z axis sin coefficients (INDATA ZAXIS_CS)

raxis_s: Any = None¶: R axis sin coefficients (lasym; RAXIS_CS)

zaxis_c: Any = None¶: Z axis cos coefficients (lasym; ZAXIS_CC)

rbc: Any = None¶: R boundary cos(m u - n nfp v)

zbs: Any = None¶: Z boundary sin(m u - n nfp v)

rbs: Any = None¶: R boundary sin (lasym)

zbc: Any = None¶: Z boundary cos (lasym)

lfreeb: bool = True¶: forced False when mgrid_file == ‘NONE’

extcur: Any = None¶: external coil-group currents [A]

nvacskip: int = 1¶: vacuum-solve cadence (<= 0 -> nfp)

classmethod from_file(path: str | Path) → VmecInput¶

Read a VMEC input file, auto-detecting INDATA vs JSON format.

Files whose first non-whitespace character is { (or with a .json suffix) are parsed as VMEC++-style JSON; everything else as a classic &INDATA Fortran namelist (VMEC2000 readin.f).

classmethod from_indata_text(text: str) → VmecInput¶: Build from &INDATA namelist text (VMEC2000 read_indata_namelist).

classmethod from_json_text(text: str) → VmecInput¶

Build from VMEC++-style JSON text (plan Appendix C / vmecpp.VmecInput).

Same key names as the dataclass fields; adiabatic_index is accepted as an alias for gamma; rbc/zbs/rbs/zbc are sparse {"m", "n", "value"} lists; axis arrays are dense. Unknown keys (e.g. VMEC++ free_boundary_method) are ignored.

to_json(path: str | Path) → Path¶

Write VMEC++-schema JSON that round-trips through from_file().

Boundary coefficients are written as sparse {"m","n","value"} lists (nonzero entries only); axis and profile arrays are dense.

to_indata(path: str | Path) → Path¶

Write a classic &INDATA namelist that round-trips exactly.

Floats are written with 17 significant digits so re-parsing reproduces the same binary values; empty arrays are omitted.

Radial profile evaluation for pressure, rotational transform, and current.

VMEC2000 counterparts: LIBSTELL/Sources/Miscel/profile_functions.f (functions pmass, piota, pcurr) as used by profil1d.f.

Every profile is a pure function of the normalized toroidal flux s in [0, 1]. All evaluation code is written with jax.numpy so the functions are usable inside jax.jit/jax.grad closures; profile kinds and knot-array shapes are static (Python-level) while coefficient values may be traced.

Units¶

pressure() returns pressure in Pascals (PRES_SCALE * pmass(x) with the VMEC input coefficients AM in Pa). VMEC2000’s pmass function returns mu0 * pres_scale * pmass (internal units, same as B**2); multiply by MU0 to obtain VMEC internal units. This matches the historical vmex.profiles.eval_profiles behavior where pressure_pa is Pa and pressure = MU0 * pressure_pa.
iota() is dimensionless (or the safety factor q = 1/iota input when lrfp=True).
current() returns VMEC’s dimensionless current shape function I(x); VMEC2000 (profil1d.f) later normalizes it so that I(1) corresponds to CURTOR when NCURR=1. That scaling is left to the caller.

Supported kinds (evaluate_profile):

kind	definition (coefficients `c` ascending, `x` in [0,1])
power_series	`sum_i c[i] x**i`
two_power	`c[0] (1 - xc[1])c[2]`
gauss_trunc	`c[0]/(1-E) * (exp(-(x/c[1])2) - E)`, `E = exp(-(1/c[1])2)` (normalized so f(0)=c[0])
pedestal	VMEC2000 `pmass` pedestal: degree-15 power series in `c[0:16]` plus a tanh pedestal shaped by `c[16:21]`
cubic_spline	VMEC `spline_cubic` through (aux_s, aux_f) knots
akima_spline	VMEC/STELLOPT `spline_akima` through the knots
line_segment	linear interpolation through the knots
power_series_ip	current: `I(x) = sum_i c[i]/(i+1) x**(i+1)` (VMEC `pcurr_type='power_series'`: c parameterizes I’)
power_series_i	current: `I(x) = sum_i c[i] x**(i+1)`
two_power_ip	current: `I(x) = int_0^x two_power(c, t) dt`
gauss_trunc_ip	current: `I(x) = int_0^x c[0](exp(-(t/c[1])**2)-E) dt`
pedestal_i	VMEC2000 `pcurr` pedestal parameterization of I(x)
cubic_spline_i(_ip)	spline of I (`_i`) or of I’ integrated (`_ip`)
akima_spline_i(_ip)	idem, Akima spline
line_segment_i(_ip)	idem, line segments

Numerical integration of the *_ip parameterized kinds uses a fixed 16-point Gauss-Legendre rule on [0, x] (VMEC2000 uses a 10-point rule; the difference is at quadrature-error level). Spline *_ip kinds are integrated analytically piecewise, as in the historical implementation.

vmex.core.profiles.MU0 = 1.2566370614359173e-06¶: Vacuum permeability [N/A^2]; VMEC2000 stel_constants uses 4e-7*pi.

vmex.core.profiles.evaluate_profile(kind: str, coefficients, aux_s, aux_f, s)¶

Evaluate one VMEC profile parameterization at s.

VMEC2000 counterpart: the SELECT CASE bodies of pmass/piota/ pcurr in profile_functions.f (one case per kind; see the module docstring for the full table).

Parameters:

kind – Profile family, case-insensitive (static under jit).
coefficients – Ascending coefficients (VMEC AM/AI/AC); ignored by the tabulated kinds.
aux_s – Knot abscissae/values (VMEC *_AUX_S/*_AUX_F); used only by the spline/line-segment kinds. Shapes must be static.
aux_f – Knot abscissae/values (VMEC *_AUX_S/*_AUX_F); used only by the spline/line-segment kinds. Shapes must be static.
s – Evaluation points; for tabulated kinds values are clipped to the knot range (VMEC evaluates only inside the knot span).

Returns:

jax.Array with the same shape as s (no pres_scale/mu0/
curtor scaling applied; see the wrapper functions).

vmex.core.profiles.pressure(pmass_type: str, am, am_aux_s, am_aux_f, s, *, pres_scale=1.0, bloat=1.0, spres_ped=1.0)¶

Pressure profile p(s) in Pascals (VMEC2000 pmass x 1/mu0).

VMEC2000 counterpart: pmass(xx) in profile_functions.f plus the spres_ped pedestal clamp applied in profil1d.f: for s > spres_ped the pressure is held at p(spres_ped).

Returns pres_scale * pmass_raw(min(|s*bloat|, 1)) in Pa. Multiply by MU0 for VMEC internal units (mu0 * Pa, the units of B**2), exactly as VMEC2000’s pmass returns mu0 * pres_scale * pmass. spres_ped is a static (host) float; pres_scale and the coefficients may be traced.

vmex.core.profiles.iota(piota_type: str, ai, ai_aux_s, ai_aux_f, s, *, bloat=1.0, lrfp=False)¶

Rotational-transform profile iota(s) (dimensionless).

VMEC2000 counterpart: piota(x) in profile_functions.f. With lrfp=True the ai coefficients parameterize the safety factor q = 1/iota and the reciprocal is returned (infinite where q = 0).

Note: VMEC2000’s piota does not apply the bloat clamp internally; the historical vmex implementation applied it uniformly and this port keeps that behavior (identical for the default bloat = 1).

vmex.core.profiles.current(pcurr_type: str, ac, ac_aux_s, ac_aux_f, s, *, bloat=1.0)¶

Enclosed-toroidal-current shape function I(s) (unnormalized).

VMEC2000 counterpart: pcurr(xx) in profile_functions.f. Kinds whose VMEC name parameterizes I' (power_series, two_power, gauss_trunc, *_ip) are integrated from 0 to x = min(|s*bloat|, 1); *_i kinds (and pedestal) parameterize I directly. VMEC2000 (profil1d.f) rescales the result so the edge value matches CURTOR when NCURR = 1; that scaling is the caller’s responsibility.

Spectral representation and physics kernels¶

Fourier-mode, angle-grid, and trigonometric-table bookkeeping for VMEC.

VMEC2000 counterparts¶

Sources/Initialization_Cleanup/fixaray.f — mode tables xm/xn, the mscale/nscale normalization, and the trig tables cosmu/sinmu/cosnv/sinnv (with derivative and integration-weighted companions).
Sources/Input_Output/read_indata.f — the internal theta grid sizes ntheta1/ntheta2/ntheta3.
Sources/Initialization_Cleanup/profil3d.f — angular integration weights (wint).

Physics/numerics conventions (parity-critical)¶

VMEC does not use a plain unweighted DFT. Its transforms rely on:

Reduced poloidal grid. For stellarator-symmetric runs (lasym=False) all angular integrals are evaluated on theta in [0, pi] (ntheta2 points, endpoints included) with half-weights at theta=0 and theta=pi. For lasym=True the full [0, 2*pi) grid (ntheta1 points, endpoint-free) is stored, but the force transforms still integrate on the reduced interval after a symmetric/antisymmetric decomposition (symforce.f).
Mode normalization. mscale(0)=nscale(0)=1 and mscale(m>=1)=nscale(n>=1)=sqrt(2). The trig tables carry these factors so that the two-stage projection (tomnsps) is the exact inverse of the two-stage synthesis (totzsps) on band-limited data (stellarator-symmetric grid), and VMEC’s internal spectral coefficients are the physical (wout) coefficients divided by mscale(m)*nscale(|n|).
Integration normalization dnorm: 1/(nzeta*(ntheta2-1)) for symmetric runs and 1/(nzeta*ntheta3) for lasym=True (SPH012314).

All tables are built once per resolution with NumPy float64 (they are static, trace-time constants for the jitted transforms in vmex.core.transforms; building them with NumPy avoids eager device transfers and is exact regardless of the JAX x64 flag).

class vmex.core.fourier.Resolution(mpol: int, ntor: int, ntheta: int, nzeta: int, nfp: int, lasym: bool, ns: int)¶

Hashable static resolution/configuration of a VMEC spectral problem.

VMEC2000: the INDATA resolution block (mpol, ntor, ntheta, nzeta, nfp, lasym, ns_array) plus the derived grid sizes from read_indata.f.

Parameters:

mpol (int) – Number of poloidal modes; m = 0 .. mpol-1.
ntor (int) – Largest toroidal mode number; n = -ntor .. ntor (n >= 0 for m=0).
ntheta (int) – Requested number of poloidal grid points (VMEC rounds it; see ntheta1/ntheta2/ntheta3).
nzeta (int) – Number of toroidal grid points per field period.
nfp (int) – Number of toroidal field periods.
lasym (bool) – True for non-stellarator-symmetric equilibria.
ns (int) – Number of radial (flux) surfaces.

property ntheta1: int¶: Full poloidal grid size, forced even (VMEC read_indata.f).

property ntheta2: int¶: Reduced poloidal grid size covering [0, pi] incl. the endpoint.

property ntheta3: int¶

ntheta1 if lasym else ntheta2.

Type:: Stored poloidal grid size

property nznt: int¶: Total angular grid size nzeta * ntheta3 (VMEC nznt).

property mnmax: int¶: Number of (m, n) modes in VMEC ordering (fixaray.f mnmax).

property lthreed: bool¶: True when the configuration is three-dimensional (ntor > 0).

class vmex.core.fourier.ModeTable(m: ndarray, n: ndarray)¶

VMEC (m, n) mode ordering plus the m-masks used downstream.

VMEC2000: the xm/xn arrays built in fixaray.f (nmin0 logic): m = 0 carries only n = 0 .. ntor; m >= 1 carries n = -ntor .. ntor. n is stored before multiplication by nfp (VMEC’s xn = n * nfp).

property mnmax: int¶: Number of (m, n) modes.

property m_is_even: ndarray¶: Even poloidal parity mask (VMEC mparity = 0).

property m_is_odd: ndarray¶: Odd poloidal parity mask (VMEC mparity = 1; carries sqrt(s)).

property m_is_zero: ndarray¶: Mask for m == 0 modes (axis/lambda special handling).

property m_is_one: ndarray¶: Mask for m == 1 modes (the residue.f90 m=1 constraint).

class vmex.core.fourier.TrigTables(ntheta1: int, ntheta2: int, ntheta3: int, nfp: int, lasym: bool, dnorm: float, dnorm3: float, mscale: ndarray, nscale: ndarray, cosmu: ndarray, sinmu: ndarray, cosmum: ndarray, sinmum: ndarray, cosmui: ndarray, sinmui: ndarray, cosmumi: ndarray, sinmumi: ndarray, cosmui3: ndarray, cosmumi3: ndarray, cosnv: ndarray, sinnv: ndarray, cosnvn: ndarray, sinnvn: ndarray, wint: ndarray)¶

Trigonometric/weight tables (VMEC2000: fixaray.f).

Theta tables have shape (ntheta3, mpol), zeta tables (nzeta, ntor+1). Suffix conventions follow fixaray.f:

cosmu/sinmu : cos(m*theta)*mscale(m) / sin(m*theta)*mscale(m).
cosmum/sinmum : theta-derivative companions (cosmum = m*cosmu, sinmum = -m*sinmu).
cosmui/sinmui : integration-weighted (dnorm and the theta-endpoint half-weights on cosmui; sinmu vanishes at 0 and pi so sinmui needs no endpoint correction).
cosmumi/sinmumi: integration-weighted derivative tables.
cosmui3/cosmumi3: full-surface-average variants (dnorm3); for lasym=False these coincide with cosmui/cosmumi.
cosnv/sinnv : cos(n*zeta)*nscale(n) / sin(n*zeta)*nscale(n).
cosnvn/sinnvn : zeta-derivative companions w.r.t. the physical toroidal angle (cosnvn = n*nfp*cosnv, sinnvn = -n*nfp*sinnv).
wint : angular integration weights on the internal grid, shape (ntheta3, nzeta) (VMEC’s wint from profil3d.f, here without the radial replication).

vmex.core.fourier.mode_table(mpol: int, ntor: int) → ModeTable¶

Build the VMEC (m, n) mode table (VMEC2000: fixaray.f, xm/xn).

Ordering: m = 0 has n = 0..ntor; m = 1..mpol-1 has n = -ntor..ntor. This matches vmex.modes.vmec_mode_table and the xm/xn arrays written to wout files.

vmex.core.fourier.trig_tables(res: Resolution) → TrigTables¶

Build the VMEC trig/weight tables for a resolution.

VMEC2000: Sources/Initialization_Cleanup/fixaray.f. The numerical conventions (mscale/nscale = sqrt(2) for nonzero mode numbers, dnorm reduced-grid normalization, exact +/-mscale values at the theta = pi endpoint, endpoint half-weights on cosmui) are ported verbatim from the parity-proven vmex.kernels.tomnsp.vmec_trig_tables.

Poloidal columns cover m = 0 .. mpol-1; toroidal columns n = 0 .. ntor.

VMEC spectral transforms: Fourier <-> real space, force projections.

VMEC2000 counterparts¶

Sources/General/totzsp_mod.f (totzsps/totzspa) — synthesis of the real-space geometry fields R, Z, lambda and their poloidal/toroidal derivatives from spectral coefficients: fourier_to_real().
Sources/General/tomnsp_mod.f (tomnsps/tomnspa) — projection of the real-space MHD force kernels back onto the Fourier basis: tomnsps() / tomnspa().
Sources/General/symforce.f — symmetric/antisymmetric decomposition of the force kernels for lasym=True runs: symforce_split().
Sources/Initialization_Cleanup/profil3d.f (scalxc) — the odd-m 1/sqrt(s) internal scaling: odd_m_sqrt_s_scaling().

Structure¶

All transforms are two-stage DFTs expressed as batched matrix products (theta stage then zeta stage, or mode-stacked phase tables), pure jax.numpy with Precision.HIGHEST, no host round-trips, and jit-friendly: the trig tables and mode tables are static (NumPy) trace-time constants.

Naming: variables use physical names with the VMEC2000 Fortran name recorded in docstrings/comments (e.g. force_R for armn, force_R_cc for frcc).

The numerics are ported verbatim from the parity-proven legacy kernels vmex/kernels/tomnsp.py (tomnsps_rzl/tomnspa_rzl) and vmex/kernels/realspace.py (vmec_realspace_synthesis* / vmec_realspace_analysis); equivalence is enforced in tests/test_fourier_transforms_ab.py.

vmex.core.transforms.register_pytree_dataclass(cls, *, meta: tuple[str, ...] = (), drop: tuple[str, ...] = ())¶

The one shared registration helper of all core result containers (R26a; formerly copied in fields/forces/geometry/residuals/setup/solver/implicit).

drop fields (which must have defaults) are excluded from the pytree entirely: they do not appear as leaves or in the treedef, and are reset to their default on unflatten. Use for host-side convenience attributes (e.g. runtime) that must not change an established pytree structure.

Fourier-space force coefficients in VMEC’s (ns, mpol, ntor+1) packing.

VMEC2000 names (tomnsp_mod.f): force_R_cc = frcc, force_R_ss = frss, force_Z_sc = fzsc, force_Z_cs = fzcs, force_lambda_sc = flsc, force_lambda_cs = flcs (symmetric, from tomnsps); force_R_sc = frsc, force_R_cs = frcs, force_Z_cc = fzcc, force_Z_ss = fzss, force_lambda_cc = flcc, force_lambda_ss = flss (antisymmetric, from tomnspa).

Suffix convention: cc multiplies cos(m theta) cos(n zeta), ss sin sin, sc sin(m theta) cos(n zeta), cs cos sin. Blocks not produced by the transform that built the object (or absent for 2D runs, ntor = 0) are None.

vmex.core.transforms.fourier_to_real(coefficient_cos: Any, coefficient_sin: Any, *, modes: ModeTable, trig: TrigTables, derivatives: Sequence[str] = ('value', 'dtheta', 'dzeta'), internal_coeffs: bool = False, odd_m_sqrt_s: bool = False, s: Any | None = None) → tuple[Any, ...]¶

Synthesize real-space fields on the VMEC internal angular grid.

VMEC2000: totzsp_mod.f — totzsps synthesizes the stellarator- symmetric blocks (rmncc/rmnss, zmnsc/zmncs, lmnsc/lmncs) and totzspa the antisymmetric ones (rmnsc/rmncs, zmncc/zmnss, lmncc/lmnss). In the signed-(m, n) mode packing used here (one cos and one sin coefficient per mode, n signed) both reduce to a single cos/sin synthesis, so this one function covers R, Z and lambda for both parities:

f(s, theta, zeta) = sum_k [ c_cos_k cos(m_k theta - n_k zeta)
                          + c_sin_k sin(m_k theta - n_k zeta) ]

Parameters:

coefficient_cos – Coefficient arrays of shape (..., ns, mnmax).
coefficient_sin – Coefficient arrays of shape (..., ns, mnmax).
derivatives – Which fields to return, from ("value", "dtheta", "dzeta"). dzeta is the derivative with respect to the physical toroidal angle (the n*nfp factor of VMEC’s xn), matching totzsps.
internal_coeffs – If True the inputs are VMEC-internal coefficients (already divided by mscale*nscale); if False (default) they are physical/wout coefficients and the conversion is applied here.
odd_m_sqrt_s – If True apply the odd-m 1/sqrt(s) factors (profil3d.f scalxc) to the coefficients before synthesis — VMEC’s totzsps consumes internal coefficients that carry this scaling.
s – Radial grid for odd_m_sqrt_s (defaults to a uniform [0, 1] grid).

Returns:

tuple of arrays, one per requested derivative, each of shape
(..., ns, ntheta3, nzeta).

vmex.core.transforms.real_to_fourier(field: Any, *, modes: ModeTable, trig: TrigTables, parity: str = 'both') → tuple[Any, Any]¶

Project a real-space field on the VMEC internal grid onto Fourier modes.

The grid-quadrature inverse of fourier_to_real() (physical/wout coefficient convention). VMEC2000 has no single subroutine for this (tomnsps projects forces, with extra kernels/masks); this is the plain analysis using the same fixaray integration weights: dnorm-normalized full-grid sums for lasym=True and the endpoint-half-weighted reduced grid [0, pi] for lasym=False.

On the reduced symmetric grid the cos and sin families are each internally orthogonal but not mutually orthogonal (cross-talk), so a definite parity should be requested for symmetric fields:

"cos": keep cos coefficients, zero the sin block (R-like fields);
"sin": keep sin coefficients, zero the cos block (Z/lambda-like);
"both": keep both (exact only on the full lasym grid).

Returns (coefficient_cos, coefficient_sin), each (ns, mnmax).

vmex.core.transforms.tomnsps(*, force_R_even: Any, force_R_odd: Any, force_R_du_even: Any, force_R_du_odd: Any, force_R_dv_even: Any, force_R_dv_odd: Any, force_Z_even: Any, force_Z_odd: Any, force_Z_du_even: Any, force_Z_du_odd: Any, force_Z_dv_even: Any, force_Z_dv_odd: Any, force_lambda_du_even: Any | None = None, force_lambda_du_odd: Any | None = None, force_lambda_dv_even: Any | None = None, force_lambda_dv_odd: Any | None = None, constraint_R_even: Any | None = None, constraint_R_odd: Any | None = None, constraint_Z_even: Any | None = None, constraint_Z_odd: Any | None = None, mpol: int, ntor: int, trig: TrigTables, include_edge: bool = False) → SpectralForce¶

Project real-space forces onto the stellarator-symmetric Fourier blocks.

VMEC2000: Sources/General/tomnsp_mod.f, tomnsps — the real-space -> Fourier-space transform of the MHD force kernels for the symmetric blocks frcc/frss (R), fzsc/fzcs (Z) and flsc/flcs (lambda).

Kernel naming (VMEC2000 names in parentheses; _even/_odd are the poloidal-parity planes, the odd one carrying the sqrt(s) representation):

force_R (armn), force_Z (azmn): multiply the basis function itself;
force_R_du (brmn), force_Z_du (bzmn), force_lambda_du (blmn): multiply the poloidal derivative of the basis;
force_R_dv (crmn), force_Z_dv (czmn), force_lambda_dv (clmn): multiply the (physical) toroidal derivative of the basis;
constraint_R (arcon), constraint_Z (azcon): the spectral- condensation constraint force, weighted here by xmpq = m(m-1).

include_edge=True keeps the boundary surface in the R/Z masks (getfsq jedge=1 semantics); the default matches fixed-boundary evolution (edge row zeroed, jmin2/jlam start indices applied).

Returns a SpectralForce with the symmetric blocks set (the ss/cs blocks are None for 2D runs, ntor = 0).

vmex.core.transforms.tomnspa(*, force_R_even: Any, force_R_odd: Any, force_R_du_even: Any, force_R_du_odd: Any, force_R_dv_even: Any, force_R_dv_odd: Any, force_Z_even: Any, force_Z_odd: Any, force_Z_du_even: Any, force_Z_du_odd: Any, force_Z_dv_even: Any, force_Z_dv_odd: Any, force_lambda_du_even: Any | None = None, force_lambda_du_odd: Any | None = None, force_lambda_dv_even: Any | None = None, force_lambda_dv_odd: Any | None = None, constraint_R_even: Any | None = None, constraint_R_odd: Any | None = None, constraint_Z_even: Any | None = None, constraint_Z_odd: Any | None = None, mpol: int, ntor: int, trig: TrigTables, include_edge: bool = False) → SpectralForce¶

Project real-space forces onto the antisymmetric Fourier blocks.

VMEC2000: Sources/General/tomnsp_mod.f, tomnspa — used for lasym=True runs after symforce_split() has separated each kernel into its symmetric and antisymmetric parts. Produces the antisymmetric blocks frsc/frcs (R), fzcc/fzss (Z) and flcc/flss (lambda).

Kernel arguments have the same meaning as in tomnsps() (pass the antisymmetric parts here). The theta/zeta stages are identical; only the output wiring differs (the roles of the cos- and sin-projected work arrays swap relative to tomnsps).

Returns a SpectralForce with the antisymmetric blocks set (the cs/ss blocks are None for 2D runs, ntor = 0).

vmex.core.transforms.symforce_split(field: Any, *, trig: TrigTables, reflect_even: bool = True) → tuple[Any, Any]¶

Split a real-space force kernel into symmetric/antisymmetric parts.

VMEC2000: Sources/General/symforce.f. For lasym=True runs, each force kernel on the full theta grid is decomposed under the stellarator reflection (theta, zeta) -> (2*pi - theta, -zeta) before the reduced- interval integrations of tomnsps() (symmetric part) and tomnspa() (antisymmetric part).

reflect_even selects the dominant symmetry of the kernel (the mapping is not uniform across kernels — see symforce.f):

True (VMEC kinds ars/bzs/bls/rcs/czs/cls, i.e. armn, bzmn, blmn, arcon, czmn, clmn): sym = (a + a_reflected)/2, asym = (a - a_reflected)/2;
False (VMEC kinds brs/azs/zcs/crs, i.e. brmn, azmn, azcon, crmn): the roles are reversed.

Rows beyond ntheta2 are kept unchanged on the symmetric output (VMEC only updates i <= ntheta2; the transforms never read those rows) and zeroed on the antisymmetric output.

Returns (field_symmetric, field_antisymmetric) with the input shape (ns, ntheta3, nzeta).

vmex.core.transforms.odd_m_sqrt_s_scaling(s: Any, mpol: int) → Any¶

VMEC odd-m 1/sqrt(s) internal scaling factors, shape (ns, mpol).

VMEC2000: profil3d.f (scalxc). Odd-m Fourier coefficients are evolved in a 1/sqrt(s) representation:

scalxc(js, m odd) = 1 / max(sqrt(s_js), sqrt(s_2))

with scalxc = 1 for even m, sqrt(s_ns) clamped to exactly 1 at the edge, and the axis clipped to the first interior half-mesh value. The same factors multiply the force coefficients after tomnsps (funct3d.f: gc = gc * scalxc), making them part of the definition of the reported residuals fsqr/fsqz/fsql.

vmex.core.transforms.physical_to_internal_scale(modes: ModeTable, trig: TrigTables) → ndarray¶

Per-mode factor converting physical (wout) to VMEC-internal coefficients.

VMEC2000: fixaray.f — internal spectral coefficients are stored divided by mscale(m) * nscale(|n|) because the trig tables carry those factors. Returns a static (mnmax,) array 1 / (mscale[m] * nscale[|n|]).

Flux-surface geometry: real-space R/Z/lambda channels and the half-mesh Jacobian.

VMEC2000 counterparts¶

Sources/General/totzsp_mod.f (totzsps/totzspa) — synthesis of the even-m / odd-m real-space geometry channels (r1, ru, rv, z1, zu, zv, lu, lv with mparity planes) including the jmin1 axis rules: real_space_geometry(), apply_lambda_axis_closure().
Sources/General/jacobian.f — half-mesh quantities r12, rs, zs, ru12, zu12, the Jacobian factor tau = sqrt(g)/R and sqrt(g) = r12 * tau, plus the Jacobian sign-change check (taumax * taumin < 0 -> irst = 2): half_mesh_jacobian().

Radial conventions (parity-critical)¶

VMEC evolves odd-m Fourier coefficients in an internal 1/sqrt(s) representation (profil3d.f scalxc), so every real-space quantity is assembled as X = X_even + sqrt(s) * X_odd where X_odd is the internal odd-m channel. Radial derivatives and the Jacobian live on the half mesh (s_half(js) = (s(js) + s(js-1)) / 2); half-mesh arrays keep the VMEC convention of copying the first interior surface into the axis slot (X(js=1) = X(js=2) in Fortran indexing).

The numerics are ported verbatim from the parity-proven legacy kernels vmex.kernels.jacobian and the geometry stage of vmex.kernels.bcovar (_compute_bcovar_parity_channels with use_vmec_synthesis=True); equivalence is enforced in tests/test_geometry_fields_ab.py.

class vmex.core.geometry.RealSpaceGeometry(R_even: Any, R_odd: Any, Z_even: Any, Z_odd: Any, dR_dtheta_even: Any, dR_dtheta_odd: Any, dZ_dtheta_even: Any, dZ_dtheta_odd: Any, dR_dzeta_even: Any, dR_dzeta_odd: Any, dZ_dzeta_even: Any, dZ_dzeta_odd: Any, dlambda_dtheta_even: Any, dlambda_dtheta_odd: Any, dlambda_dzeta_even: Any, dlambda_dzeta_odd: Any)¶

Even-m / odd-m real-space geometry channels on the internal grid.

VMEC2000 names (totzsp_mod.f outputs; _even/_odd are the mparity = 0/1 planes, odd carrying the internal 1/sqrt(s) representation): R_even/R_odd = r1, dR_dtheta_* = ru, dR_dzeta_* = rv, Z_* = z1, dZ_dtheta_* = zu, dZ_dzeta_* = zv, dlambda_dtheta_* = lu and dlambda_dzeta_* (note VMEC stores lv = -d(lambda)/dzeta; here the plain derivative is stored and the sign lives in vmex.core.fields.contravariant_b()).

The odd planes carry the jmin1 axis rules of totzsps: on the axis row only the m = 1 content survives, extrapolated from the first interior surface (jmin1(m=1) = 1, jmin1(m>=2) = 2).

All arrays have shape (ns, ntheta3, nzeta).

theta_derivatives_full(s: Any) → tuple[Any, Any]¶

Return full-mesh (ru0, zu0) = X_even + sqrt(s) * X_odd.

VMEC2000: ru0/zu0 formed in funct3d.f and consumed by the constraint scaling (bcovar.f arnorm/aznorm) and the constraint force (alias.f ztemp).

class vmex.core.geometry.HalfMeshJacobian(r12: Any, dR_ds: Any, dZ_ds: Any, ru12: Any, zu12: Any, tau: Any, sqrt_g: Any, jacobian_sign_changed: Any)¶

Half-mesh Jacobian quantities (VMEC2000: Sources/General/jacobian.f).

Attributes (VMEC names in parentheses; all (ns, ntheta3, nzeta) except the flag):

r12 (r12): R interpolated to the half mesh.
dR_ds (rs), dZ_ds (zs): radial finite differences on the half mesh, including the odd-m sqrt(s) chain-rule terms.
ru12 (ru12), zu12 (zu12): poloidal derivatives averaged to the half mesh.
tau (tau): sqrt(g)/R = ru12*zs - rs*zu12 + (odd-m d(sqrt s)/ds corrections) — eq. (10) of Hirshman & Whitson with the discrete half-mesh corrections of jacobian.f.
sqrt_g (gsqrt): the coordinate Jacobian r12 * tau.
jacobian_sign_changed: scalar bool, True when tau changes sign over the interior surfaces (taumax*taumin < 0 -> VMEC irst = 2 / bad_jacobian_flag). Computed without host branching.

vmex.core.geometry.apply_lambda_axis_closure(lambda_sin: Any, *, modes: ModeTable, ntor: int) → Any¶

Apply VMEC’s symmetric-lambda axis closure to the sin coefficients.

VMEC2000: totzsp_mod.f (totzsps) sets lmncs(js=1, n>0, m=0) = lmncs(js=2, n>0, m=0) for three-dimensional runs (lthreed). In the signed-(m, n) packing used here this copies the first interior surface of the (m=0, n>0) sin coefficients into the axis row. A no-op for 2D runs (ntor = 0) or ns < 2.

vmex.core.geometry.real_space_geometry(*, R_cos: Any, R_sin: Any, Z_cos: Any, Z_sin: Any, lambda_cos: Any, lambda_sin: Any, modes: ModeTable, trig: TrigTables, s: Any) → RealSpaceGeometry¶

Synthesize the VMEC even-m/odd-m geometry channels from coefficients.

VMEC2000: totzsp_mod.f — totzsps (+ totzspa content in the signed-(m, n) packing) producing r1, ru, rv, z1, zu, zv, lu, lv split by poloidal-mode parity, with the profil3d.f scalxc odd-m 1/sqrt(s) scaling and the jmin1 axis rules applied.

Parameters:

R_cos – Spectral coefficients, shape (ns, mnmax), in VMEC internal normalization (divided by mscale*nscale) and in the physical basis (the residue.f90 m=1 constraint already undone; see vmex.kernels.parity.vmec_m1_internal_to_physical_signed). lambda_sin should already carry the 3D axis closure (apply_lambda_axis_closure()).
R_sin – Spectral coefficients, shape (ns, mnmax), in VMEC internal normalization (divided by mscale*nscale) and in the physical basis (the residue.f90 m=1 constraint already undone; see vmex.kernels.parity.vmec_m1_internal_to_physical_signed). lambda_sin should already carry the 3D axis closure (apply_lambda_axis_closure()).
Z_cos – Spectral coefficients, shape (ns, mnmax), in VMEC internal normalization (divided by mscale*nscale) and in the physical basis (the residue.f90 m=1 constraint already undone; see vmex.kernels.parity.vmec_m1_internal_to_physical_signed). lambda_sin should already carry the 3D axis closure (apply_lambda_axis_closure()).
Z_sin – Spectral coefficients, shape (ns, mnmax), in VMEC internal normalization (divided by mscale*nscale) and in the physical basis (the residue.f90 m=1 constraint already undone; see vmex.kernels.parity.vmec_m1_internal_to_physical_signed). lambda_sin should already carry the 3D axis closure (apply_lambda_axis_closure()).
lambda_cos – Spectral coefficients, shape (ns, mnmax), in VMEC internal normalization (divided by mscale*nscale) and in the physical basis (the residue.f90 m=1 constraint already undone; see vmex.kernels.parity.vmec_m1_internal_to_physical_signed). lambda_sin should already carry the 3D axis closure (apply_lambda_axis_closure()).
lambda_sin – Spectral coefficients, shape (ns, mnmax), in VMEC internal normalization (divided by mscale*nscale) and in the physical basis (the residue.f90 m=1 constraint already undone; see vmex.kernels.parity.vmec_m1_internal_to_physical_signed). lambda_sin should already carry the 3D axis closure (apply_lambda_axis_closure()).
s – Full-mesh radial grid, shape (ns,) (uniform, s in [0, 1]).

Returns:

RealSpaceGeometry with all channels on the internal
(ns, ntheta3, nzeta) grid.

vmex.core.geometry.half_mesh_jacobian(geometry: RealSpaceGeometry, *, s: Any) → HalfMeshJacobian¶

Half-mesh Jacobian from the even/odd geometry channels.

VMEC2000: Sources/General/jacobian.f — with X = X_even + sqrt(s) * X_odd on the full mesh, builds (half-mesh index js; Fortran js = 2..ns):

r12 = 0.5*(R_e(js) + R_e(js-1) + shalf(js)*(R_o(js) + R_o(js-1)))
rs = ohs*(R_e(js) - R_e(js-1) + shalf(js)*(R_o(js) - R_o(js-1))) (zs analogous, ru12/zu12 are the poloidal-derivative averages)
tau = ru12*zs - rs*zu12 + 0.25*[ (ru_o*z1_o + ...) + (ru_e*z1_o + ...)/shalf ] — the d(sqrt s)/ds correction terms
sqrt_g = r12 * tau (VMEC: gsqrt), zeroed on the axis row.

The axis row of each half-mesh array copies the first interior surface (VMEC serial convention). jacobian_sign_changed reproduces the jacobian.f check taumax*taumin < 0 (over js = 2..ns) that sets irst = 2 (bad_jacobian_flag); it is returned as a traced boolean — reacting to it is the solver’s job (no host branching here).

vmex.core.geometry.sqrt_s_half_mesh(s: Any) → Any¶

sqrt(s) on the VMEC radial half mesh (VMEC: pshalf/shalf).

VMEC2000: profil3d.f — shalf(js) = sqrt(hs * (js - 1.5)). The axis slot repeats the first interior half-mesh value (VMEC’s arrays are simply never read there, but legacy kernels fill it this way and downstream formulas rely on it). Shape (ns,) in, (ns,) out.

Magnetic fields, metric elements, energies, force norms and constraint scaling.

VMEC2000 counterparts¶

All functions in this module port pieces of Sources/General/bcovar.f (and the helpers it calls):

metric_elements() — half-mesh guu/guv/gvv (bcovar.f metric block).
lambda_scale() — lamscale (profil1d.f).
magnetic_fields() — contravariant B^u/B^v (bcovar.f + add_fluxes.f), covariant B_u/B_v, differential volume vp, the mass -> pressure closure pres = mass / vp**gamma and the total pressure bsq.
energies_and_force_norms() — energies wb/wp, plasma volume and the residual normalizations fnorm/fnormL + r1 (bcovar.f, getfsq).
preconditioned_force_norm() — fnorm1 (bcovar.f).
surface_currents() — buco/bvco (fbal.f) and the derived scalars ctor, rbtor, rbtor0 (bcovar.f).
constraint_scaling() — the spectral-condensation strength tcon(js) (bcovar.f + the diagonal elements of precondn.f).

All functions are pure and jit-friendly (no host round-trips, no value-based Python branching). The VMEC2000 recompute cadence (ns4 = 25 iterations for the preconditioner/norms/tcon refresh) is deliberately not implemented here — caching is the solver’s job.

The numerics are ported verbatim from the parity-proven legacy kernels vmex.kernels.bcovar, vmex.kernels.residue, vmex.kernels.lforbal and vmex.kernels.constraints; equivalence is enforced in tests/test_geometry_fields_ab.py.

class vmex.core.fields.MetricElements(guu: Any, guv: Any, gvv: Any)¶

Half-mesh metric elements (VMEC: guu, guv, gvv in bcovar.f).

guu = R_u^2 + Z_u^2, guv = R_u R_v + Z_u Z_v, gvv = R_v^2 + Z_v^2 + R^2, each assembled from the even/odd-m channels and interpolated to the radial half mesh; guv/gvv are zeroed on the axis row. Shapes (ns, ntheta3, nzeta).

class vmex.core.fields.MagneticFields(bsupu: Any, bsupv: Any, bsubu: Any, bsubv: Any, total_pressure: Any, pressure: Any, vp: Any, lamscale: Any, chips: Any)¶

Half-mesh magnetic field state (VMEC2000: bcovar.f).

bsupu/bsupv (VMEC bsupu, bsupv): contravariant components B^u = phipog*(chip + lamscale*d(lambda)/dzeta) and B^v = phipog*(phip + lamscale*d(lambda)/dtheta) with phipog = 1/sqrt(g) on the half mesh (zeroed on axis).
bsubu/bsubv (bsubu, bsubv): covariant components B_u = guu B^u + guv B^v, B_v = guv B^u + gvv B^v.
total_pressure (bsq): |B|^2/2 + p with |B|^2 = B^u B_u + B^v B_v.
pressure (pres, shape (ns,)): kinetic pressure on the half mesh, mass/vp**gamma when a mass profile is given (adiabatic closure), zero on axis.
vp (vp): differential volume dV/ds / (2 pi)^2 = signgs * <sqrt(g)> per half-mesh surface, shape (ns,).
lamscale: the lambda normalization (profil1d.f).
chips (chips): the effective half-mesh d(chi)/ds actually used in B^u — equal to the input except in the ncurr = 1 current-constrained mode (add_fluxes.f), shape (ns,).

class vmex.core.fields.EnergiesAndForceNorms(wb: Any, wp: Any, volume: Any, vp: Any, energy_density: Any, fnorm: Any, fnormL: Any, r1: Any)¶

Energies and force-residual normalizations (VMEC2000: bcovar.f).

wb (wb): magnetic energy |hs * sum_{js>=2} <signgs*sqrt(g)* |B|^2/2>| (the wout wb normalization; multiply by (2 pi)^2 for Joules/mu0).
wp (wp): kinetic energy hs * sum_{js>=2} vp*pres.
volume: plasma volume / (2 pi)^2 = hs * sum(vp(2:ns)).
vp: per-surface differential volume (same as MagneticFields).
energy_density (r2): max(wb, wp)/volume.
fnorm (fnorm): 1 / (sum(guu*r12^2*wint) * r2^2) — normalizes the physical R/Z force residuals fsqr/fsqz (getfsq).
fnormL (fnormL): 1 / (sum((B_u^2 + B_v^2)*wint) * lamscale^2) — normalizes the lambda residual fsql.
r1: the 1/(2*r0scale)^2 = 1/4 prefactor applied with fnorm in getfsq.

class vmex.core.fields.CurrentDiagnostics(buco: Any, bvco: Any, ctor: Any, rbtor: Any, rbtor0: Any)¶

Flux-surface current averages and derived scalars (VMEC2000 names).

buco (buco): <B_u> per half-mesh surface — the enclosed toroidal current profile (fbal.f), shape (ns,).
bvco (bvco): <B_v> — the poloidal current profile.
ctor (ctor): net toroidal plasma current (mu0*I), signgs * 2 pi * (1.5*buco(ns) - 0.5*buco(ns-1)) (bcovar.f).
rbtor (rbtor): R*Btor at the edge, 1.5*bvco(ns) - 0.5*bvco(ns-1).
rbtor0 (rbtor0): R*Btor on axis, 1.5*bvco(2) - 0.5*bvco(3).

vmex.core.fields.metric_elements(geometry: RealSpaceGeometry, *, s: Any) → MetricElements¶

Half-mesh metric elements guu, guv, gvv.

VMEC2000: bcovar.f — with X = X_even + sqrt(s)*X_odd the squared / cross products decompose into even parts (a0^2 + b0^2 + s*(a1^2 + b1^2)) and odd parts (2*(a0*a1 + b0*b1)) which are then averaged to the half mesh with the sqrt(s_half) weight on the odd piece:

guu = R_u^2 + Z_u^2
guv = R_u R_v + Z_u Z_v
gvv = R_v^2 + Z_v^2 + R^2 (the R^2 term is the cylindrical toroidal metric; VMEC works at unit d(zeta_physical)/d(zeta))

guv and gvv are zeroed on the axis row (VMEC never reads them there and the legacy kernels enforce it).

vmex.core.fields.lambda_scale(phips: Any, s: Any) → Any¶

VMEC lamscale = sqrt(hs * sum_{js=2..ns} phips(js)^2) (profil1d.f).

phips is the half-mesh d(phi)/ds / (2 pi) profile (VMEC internal units). Lambda is evolved as lamscale * lambda internally; the factor reappears in B^u/B^v and in fnormL.

vmex.core.fields.magnetic_fields(*, geometry: RealSpaceGeometry, jacobian: HalfMeshJacobian, metrics: MetricElements, trig: TrigTables, s: Any, phips: Any, phipf: Any, chips: Any, signgs: int, gamma: float = 0.0, pressure: Any | None = None, mass: Any | None = None, ncurr: int = 0, enclosed_current: Any | None = None) → MagneticFields¶

Contravariant/covariant B, differential volume and total pressure.

VMEC2000: bcovar.f (field/pressure blocks) + add_fluxes.f:

lu = lamscale * d(lambda)/dtheta + phip, lv = -lamscale * d(lambda)/dzeta on the full mesh (even/odd planes);
half-mesh staggering: B^v = phipog * <lu>, B^u = phipog * <lv> + chip * phipog with phipog = 1/sqrt(g) (overg), both zero on the axis row;
ncurr = 1 (add_fluxes.f): the half-mesh chips is solved from the prescribed enclosed toroidal current profile icurv: chips = (icurv - <guu B^u_lam + guv B^v>) / <guu/sqrt(g)>;
B_u = guu B^u + guv B^v, B_v = guv B^u + gvv B^v;
vp = signgs * <sqrt(g)> per surface (dV/ds/(2 pi)^2);
adiabatic pressure closure pres = mass / vp**gamma (zero where vp = 0, i.e. on axis) when mass is given, else the supplied half-mesh pressure profile is used;
bsq = |B|^2/2 + pres (VMEC bsq, the total pressure).

Parameters:

phips – VMEC-internal flux derivative profiles, shape (ns,): half-mesh phips, full-mesh phip (= signgs*phipf_wout/(2 pi)), and half-mesh chips (d(chi)/ds).
phipf – VMEC-internal flux derivative profiles, shape (ns,): half-mesh phips, full-mesh phip (= signgs*phipf_wout/(2 pi)), and half-mesh chips (d(chi)/ds).
chips – VMEC-internal flux derivative profiles, shape (ns,): half-mesh phips, full-mesh phip (= signgs*phipf_wout/(2 pi)), and half-mesh chips (d(chi)/ds).
signgs – Jacobian orientation sign (+1/-1).
gamma – Adiabatic index (INDATA GAMMA); used only with mass.
pressure – Half-mesh profiles (ns,) in VMEC internal units (mu0*Pa). Exactly one should be provided; mass takes precedence.
mass – Half-mesh profiles (ns,) in VMEC internal units (mu0*Pa). Exactly one should be provided; mass takes precedence.
ncurr – ncurr = 1 activates the current-constrained mode with the prescribed icurv profile (enclosed_current, defaults to zero).
enclosed_current – ncurr = 1 activates the current-constrained mode with the prescribed icurv profile (enclosed_current, defaults to zero).

vmex.core.fields.energies_and_force_norms(*, jacobian: HalfMeshJacobian, metrics: MetricElements, fields: MagneticFields, trig: TrigTables, s: Any, signgs: int) → EnergiesAndForceNorms¶

Energies wb/wp and residual normalizations fnorm/fnormL.

VMEC2000: bcovar.f —

wb = |hs * sum_{js=2..ns} <signgs*sqrt(g) * |B|^2/2>|
wp = hs * sum_{js=2..ns} vp*pres
volume = hs * sum(vp(2:ns)), r2 = max(wb, wp)/volume
fnorm = 1 / (sum(guu * r12^2 * wint) * r2^2) (R/Z force norm)
fnormL = 1 / (sum((B_u^2 + B_v^2) * wint) * lamscale^2) (lambda)
r1 = 1/(2*r0scale)^2 — the constant companion factor in getfsq (fsqr = r1 * fnorm * |F_R|^2 etc.).

For exact legacy parity the surface pressure entering wp is recovered from bsq - |B|^2/2 (numerically identical to using fields.pressure directly up to roundoff relative to wb).

vmex.core.fields.preconditioned_force_norm(*, R_cos: Any, Z_sin: Any, modes: ModeTable, R_sin: Any | None = None, Z_cos: Any | None = None) → Any¶

Preconditioned R/Z force normalization fnorm1.

VMEC2000: bcovar.f — fnorm1 = 1/sum(xc**2) over the internal R/Z spectral state. In the signed-(m, n) packing this is

fnorm1 = 1 / sum_{js>=2} w_k * (Rcos_k^2 + Zsin_k^2 [+ Rsin_k^2 + Zcos_k^2 for lasym])

with w_k = 2 for helical modes (m > 0 and n != 0; the internal cc/ss blocks each carry half the signed content) and w_k = 1 otherwise, and the (m, n) = (0, 0) mode excluded from the Rcos block (VMEC skips the R00 profile). The axis surface is excluded (bcovar_par accumulates over l = 2..ns). Ported from the parity-proven ModeTransform.rz_norm in the legacy solver.

Pass the internal (evolved) coefficients — the m=1 constrained representation — exactly as VMEC does with xc.

vmex.core.fields.surface_currents(*, bsubu: Any, bsubv: Any, trig: TrigTables, s: Any, signgs: int) → CurrentDiagnostics¶

Flux-surface current averages buco/bvco and ctor/rbtor/rbtor0.

VMEC2000: fbal.f (buco = <B_u>, bvco = <B_v> with the axis-masked angular weights) and bcovar.f for the edge/axis extrapolations:

ctor = signgs * 2 pi * (1.5*buco(ns) - 0.5*buco(ns-1)) — the net toroidal plasma current in mu0*Ampere units;
rbtor = 1.5*bvco(ns) - 0.5*bvco(ns-1) (edge R*Btor, VMEC fpsi(ns) extrapolation);
rbtor0 = 1.5*bvco(2) - 0.5*bvco(3) (axis R*Btor).

vmex.core.fields.constraint_scaling(*, tcon0: Any, geometry: RealSpaceGeometry, jacobian: HalfMeshJacobian, total_pressure: Any, trig: TrigTables, s: Any) → Any¶

Spectral-condensation constraint strength tcon(js).

VMEC2000: bcovar.f (with the m = 0 diagonal elements of precondn.f) — per plan Appendix D:

tcon(js) = min(|ard(js,1)|/arnorm(js), |azd(js,1)|/aznorm(js))
           * tcon_multiplier * (32*hs)^2    for js = 2..ns-1,
tcon(ns) = 0.5 * tcon(ns-1),

where (precondn.f, with ptau = pfactor*r12^2*bsq*wint/sqrt(g) and pfactor = -4*r0scale^2):

ax(js,1) = sum(ptau * (zu12*ohs)^2), ard(js,1) = ax(js,1) + ax(js+1,1) (the Z-derivative pair feeds the R diagonal, and vice versa for azd);
arnorm(js) = <ru0^2>, aznorm(js) = <zu0^2> with ru0 = ru_even + sqrt(s)*ru_odd;
tcon_multiplier = min(|tcon0|, 1)*(1 + ns*(1/60 + ns/(200*120)))/(4*r0scale^2)^2 (the resolution-dependent ramp of bcovar.f).

The axis slot holds the clamped tcon0 (VMEC initializes tcon(:) = tcon0 before overwriting the interior; the value is never used by the constraint operator). The ns4 = 25-iteration refresh cadence of VMEC2000 lives in the solver, not here — this is the pure recompute.

Real-space MHD force kernels and the spectral-condensation constraint force.

VMEC2000 counterparts¶

Sources/General/forces.f — the cylindrical R/Z force kernels armn/brmn/crmn and azmn/bzmn/czmn (even/odd-m planes) built from the half-mesh field/geometry quantities of bcovar.f: mhd_force_kernels().
Sources/General/bcovar.f (lambda-force block) — the full-mesh covariant lambda force kernels blmn/clmn from bsubu/bsubv: lambda_force_kernels().
Sources/General/alias.f + funct3d.f (constraint block) — the spectral-condensation constraint force gcon from ztemp = (rcon - rcon0)*ru0 + (zcon - zcon0)*zu0: constraint_force() / alias_constraint_force(), with the faccon(m) weights of fixaray.f (faccon()).
Sources/General/symforce.f — per-kernel symmetric/antisymmetric split for lasym runs (delegating to vmex.core.transforms.symforce_split()): symmetrize_forces().

Equations (Hirshman & Whitson 1983; VMEC2000 forces.f)¶

With lu = bsq*R12 and lv = bsq*tau on the half mesh (bsq = |B|^2/2 + p) and the “GIJ” products guu = B^u B^u sqrt(g), guv = B^u B^v sqrt(g), gvv = B^v B^v sqrt(g), the spectral force components are

F_R = armn - d(brmn)/dtheta + d(crmn)/dzeta F_Z = azmn - d(bzmn)/dtheta + d(czmn)/dzeta F_lambda = -d(blmn)/dtheta + d(clmn)/dzeta

where the A kernels carry the radial derivative of the MHD energy (ohs-differenced half-mesh terms) plus the R-curvature term -gvv*R, the B kernels the poloidal metric terms, and the C kernels the toroidal metric terms. Odd-m planes carry the internal sqrt(s) representation and the d(sqrt s)/ds = 1/(2 sqrt s) chain-rule terms (dshalfds = 0.25 discrete factor).

The numerics are ported verbatim from the parity-proven legacy kernels vmex.kernels.forces (_assemble_vmec_rz_radial_forces, _constraint_kernels_from_state), vmex.kernels.bcovar (_compute_bcovar_lambda_force_assembly) and vmex.kernels.constraints (alias_gcon/faccon_from_signgs); equivalence is enforced in tests/test_forces_residuals_ab.py. All functions are pure jax.numpy (jit-friendly, no host round-trips).

class vmex.core.forces.RealSpaceForces(force_R_even: Any, force_R_odd: Any, force_R_du_even: Any, force_R_du_odd: Any, force_R_dv_even: Any, force_R_dv_odd: Any, force_Z_even: Any, force_Z_odd: Any, force_Z_du_even: Any, force_Z_du_odd: Any, force_Z_dv_even: Any, force_Z_dv_odd: Any, force_lambda_du_even: Any, force_lambda_du_odd: Any, force_lambda_dv_even: Any, force_lambda_dv_odd: Any, constraint_R_even: Any, constraint_R_odd: Any, constraint_Z_even: Any, constraint_Z_odd: Any, gcon: Any | None = None, rcon0: Any | None = None, zcon0: Any | None = None)¶

Real-space MHD force kernels on the internal grid, split by m-parity.

VMEC2000 names (forces.f/bcovar.f; _even/_odd are the mparity = 0/1 planes, odd carrying the internal sqrt(s) representation): force_R = armn, force_R_du = brmn, force_R_dv = crmn, force_Z = azmn, force_Z_du = bzmn, force_Z_dv = czmn, force_lambda_du = blmn, force_lambda_dv = clmn, constraint_R = arcon, constraint_Z = azcon. Field names match the keyword arguments of vmex.core.transforms.tomnsps() so the projection call is direct.

gcon is the spectral-condensation constraint force field (alias.f); rcon0/zcon0 are the constraint baselines actually used (funct3d.f — the free-boundary x0.9 per-iteration ramp is the solver’s job). All arrays have shape (ns, ntheta3, nzeta).

vmex.core.forces.faccon(mpol: int, signgs: int) → ndarray¶

Constraint-force mode weights faccon(m) (VMEC2000 fixaray.f).

faccon(0) = faccon(mpol-1) = 0 and, for m = 1 .. mpol-2,

faccon(m) = -0.25 * signgs / xmpq(m+1, 1)**2

with xmpq(m, 1) = m*(m-1), i.e. the constraint spectrum is restricted to m in [1, mpol-2] (Appendix D). Static (NumPy) output.

vmex.core.forces.lambda_force_kernels(*, geometry: RealSpaceGeometry, jacobian: HalfMeshJacobian, metrics: MetricElements, fields: MagneticFields, s: Any, phipf: Any) → tuple[Any, Any, Any, Any]¶

Full-mesh covariant lambda force kernels (VMEC2000 bcovar.f).

The lambda force needs bsubu/bsubv on the radial full mesh. bsubu is the plain forward average of the half-mesh covariant component; bsubv is reconstructed directly on the full mesh from the lambda derivatives (lvv = gvv/sqrt(g), bsubv = <lvv>*lu_even + <lvv*shalf>*lu_odd + <guv*B^u>) and then blended with the plain average using the v8.49 damping bdamp = 2*0.05*(1 - s):

bsubv_full = bdamp*bsubv_reconstructed + (1 - bdamp)*<bsubv>.

Both are scaled by -lamscale on js >= 2 (the axis row is inactive, jlam = 2) and the odd planes carry the extra sqrt(s) factor:

clmn = -lamscale * bsubu_full (multiplies d(basis)/dzeta), blmn = -lamscale * bsubv_full (multiplies d(basis)/dtheta).

Parameters mirror the pipeline objects; phipf is the full-mesh VMEC-internal phip profile (the same one passed to vmex.core.fields.magnetic_fields()).

Returns (blmn_even, blmn_odd, clmn_even, clmn_odd) — the force_lambda_du_even/odd and force_lambda_dv_even/odd kernels.

vmex.core.forces.mhd_force_kernels(*, geometry: RealSpaceGeometry, jacobian: HalfMeshJacobian, fields: MagneticFields, s: Any, lthreed: bool) → tuple[Any, ...]¶

The twelve cylindrical R/Z force kernels (VMEC2000 forces.f).

Inputs are the half-mesh quantities that bcovar.f leaves behind for forces: lu = bsq*R12, lv = bsq*tau and the “GIJ” B-products guu/guv/gvv = B^i B^j sqrt(g) (all with the axis row zeroed), plus the full-mesh even/odd geometry channels. The radial staggering (forward differences/averages/sums along js with the dshalfds = 0.25 odd-m chain-rule terms) is ported verbatim from forces.f.

For 2D runs (lthreed = False) the crmn/czmn kernels are never consumed by the zeta stage of tomnsps; the returned placeholders reproduce VMEC’s storage-overlay values (crmn = lv*shalf, czmn = lu, crmn_odd = -lamscale*dlambda_dzeta_odd, czmn_odd = lu_odd) so kernel-level parity holds bit-for-bit.

Returns the tuple (force_R_even, force_R_odd, force_R_du_even, force_R_du_odd, force_R_dv_even, force_R_dv_odd, force_Z_even, force_Z_odd, force_Z_du_even, force_Z_du_odd, force_Z_dv_even, force_Z_dv_odd) (VMEC armn/brmn/crmn/azmn/bzmn/czmn).

vmex.core.forces.alias_constraint_force(ztemp: Any, *, trig: TrigTables, mpol: int, ntor: int, signgs: int, tcon: Any, lasym: bool) → Any¶

De-aliased constraint force gcon from ztemp (VMEC2000 alias.f).

Forward-transforms ztemp on the reduced theta interval, keeps only the m in [1, mpol-2] band (via the faccon() weights, which vanish at m = 0 and m = mpol-1), scales by tcon(js), and inverse transforms. For lasym=True the alias_par reflection algebra is used: symmetric/antisymmetric zeta blocks are built from ztemp and its stellarator reflection, and the second theta half-interval is filled by gcon(2*pi - theta, -zeta) = -gcon_s + gcon_a.

ztemp has shape (ns, ntheta3, nzeta); the returned gcon matches it.

vmex.core.forces.constraint_force(*, R_cos: Any, R_sin: Any, Z_cos: Any, Z_sin: Any, geometry: RealSpaceGeometry, modes: ModeTable, trig: TrigTables, s: Any, tcon: Any, signgs: int, rcon0: Any | None = None, zcon0: Any | None = None) → tuple[Any, Any, Any, Any, Any]¶

Spectral-condensation constraint force pipeline (funct3d.f + alias.f).

Synthesize rcon/zcon: the real-space image of the xmpq(m,1) = m*(m-1)-weighted R/Z coefficients (physical-signed internal normalization, same input convention as vmex.core.geometry.real_space_geometry()), assembled as X_even + sqrt(s) * X_odd_internal with the jmin1 axis rules.
Baselines: rcon0/zcon0 default to the edge profile scaled by s (funct3d.f fixed-boundary initialization rcon0 = s * rcon(ns)); pass persisted arrays to keep VMEC’s caching/free-boundary 0.9-ramp semantics (that ramp lives in the solver).
ztemp = (rcon - rcon0)*ru0 + (zcon - zcon0)*zu0 with the physical full-mesh theta derivatives ru0/zu0.
gcon via alias_constraint_force() with tcon(js) scaling.

Returns (gcon, rcon, zcon, rcon0, zcon0).

vmex.core.forces.mhd_forces(*, geometry: RealSpaceGeometry, jacobian: HalfMeshJacobian, metrics: MetricElements, fields: MagneticFields, R_cos: Any, R_sin: Any, Z_cos: Any, Z_sin: Any, modes: ModeTable, trig: TrigTables, s: Any, phipf: Any, tcon: Any, signgs: int, rcon0: Any | None = None, zcon0: Any | None = None) → RealSpaceForces¶

Assemble all real-space force kernels (funct3d.f force stage).

Combines mhd_force_kernels(), lambda_force_kernels() and constraint_force(), applying the forces.f constraint add-back

brmn += (rcon - rcon0)*gcon, bzmn += (zcon - zcon0)*gcon

(odd planes with the extra sqrt(s)), and forming the constraint projection kernels arcon = ru0*gcon / azcon = zu0*gcon consumed by tomnsps with the xmpq weights.

The R/Z/lambda coefficients are the physical-signed VMEC-internal ones used for vmex.core.geometry.real_space_geometry() (the m=1 constraint of residue.f90 already undone — see vmex.core.residuals.m1_constrained_to_physical()). Passing tcon = 0 disables the constraint force without retracing.

vmex.core.forces.symmetrize_forces(forces: RealSpaceForces, *, trig: TrigTables) → tuple[RealSpaceForces, RealSpaceForces]¶

Symmetric/antisymmetric decomposition of every kernel (symforce.f).

For lasym=True runs each kernel is split under the stellarator reflection (theta, zeta) -> (2*pi - theta, -zeta) before the reduced interval projections; the dominant symmetry differs per kernel (_SYMFORCE_REFLECT_EVEN, matching the symforce.f kinds). Returns (forces_symmetric, forces_antisymmetric); the gcon/rcon0/zcon0 diagnostics are carried on the symmetric output only.

vmex.core.forces.spectral_mhd_forces(forces: RealSpaceForces, *, mpol: int, ntor: int, trig: TrigTables, include_edge: bool = False) → SpectralForce¶

Project the real-space force kernels onto the Fourier basis.

VMEC2000: funct3d.f — tomnsps for stellarator-symmetric runs; for lasym=True the kernels are first decomposed with symmetrize_forces() (symforce.f) and the antisymmetric parts go through tomnspa. Returns a single vmex.core.transforms.SpectralForce carrying both block families (antisymmetric blocks are None for symmetric runs).

Force residual scalars, the m=1 polar constraint and preconditioned residuals.

VMEC2000 counterparts¶

Sources/General/residue.f90 — the m=1 polar constraint: coefficient-space mappings m1_constrained_to_physical() / m1_physical_to_constrained() (readin.f lconm1 convention, rbss = (rbs + zbc)/2), the force-side rotation (gcr, gcz)_{m=1} -> ((gcr+gcz)/sqrt2, (gcr-gcz)/sqrt2) (m1_residue_rotation(), constrain_m1), its conditional release (zero_m1_z_force() + m1_zero_condition()), and the scale_m1 preconditioner factors (scale_m1_preconditioner_rhs()).
Sources/General/getfsq.f — the invariant residuals fsqr/fsqz (gnorm = r1*fnorm, jsmax = ns-1 + medge) plus fsql and the edge diagnostic fedge: force_residuals() (+ edge_force_condition() for the free-boundary medge = 1 rule).
Sources/TimeStep/funct3d.f — the post-tomnsps gc = gc*scalxc odd-m scaling (scalxc_scale_force()) and the preconditioned residuals fsqr1/fsqz1 = fnorm1*sum(gc**2), fsql1 = hs*sum(gcl_preconditioned**2) (preconditioned_residuals(), with apply_radial_preconditioner() / apply_lambda_preconditioner() wrapping vmex.core.preconditioner).

Conventions (parity-critical)¶

The m=1 constraint makes the poloidal angle unique near the axis: VMEC evolves rss_int = (rss + zcs)/2 and zcs_int = (rss - zcs)/2 (and rsc/zcc for lasym) instead of the physical pairs. Geometry synthesis therefore needs m1_constrained_to_physical() first (this is the mapping assumed by vmex.core.geometry.real_space_geometry()). All release conditions are pure functions of traced values — the solver supplies the previous fsqz/iteration counters and receives traced booleans, so a single compiled step serves both execution lanes.

The numerics are ported verbatim from the parity-proven legacy kernels vmex.kernels.residue, vmex.kernels.parity and the solver modules solvers/fixed_boundary/residual/payload_blocks.py / preconditioning/operators.py; equivalence is enforced in tests/test_forces_residuals_ab.py.

vmex.core.residuals.M1_FSQZ_RELEASE_THRESHOLD = 1e-06¶

the constrained m=1 Z force is zeroed once fsqz drops below this (or during the first iterations after a restart).

Type:: residue.f90 FThreshold

vmex.core.residuals.M1_STARTUP_ITERATIONS = 2¶

iter2 < 2 also zeroes the constrained Z force (reset support).

Type:: v8.50

vmex.core.residuals.EDGE_FORCE_FSQ_THRESHOLD = 1e-06¶

free-boundary edge rows join getfsq (jedge = 1) when fsqr + fsqz has dropped below this …

Type:: residue.f90 (SPH041117)

vmex.core.residuals.EDGE_FORCE_ITERATION_WINDOW = 50¶: … within this many iterations of the last restart (delIter < 50).

class vmex.core.residuals.ForceResiduals(fsqr: Any, fsqz: Any, fsql: Any, fedge: Any, gcr2: Any, gcz2: Any, gcl2: Any)¶

Invariant force residuals (VMEC2000 getfsq.f / residue.f90).

fsqr/fsqz (fsqr, fsqz): r1*fnorm * sum(gc**2) over the evolved surfaces (jsmax = ns-1 + medge);
fsql (fsql): fnormL * sum(gcl**2) (all surfaces);
fedge (fedge): r1*fnorm times the edge-row R/Z sums — only meaningful when the projection kept the edge row (include_edge);
gcr2/gcz2/gcl2: the raw sums of squares (solver diagnostics).

class vmex.core.residuals.PreconditionedResiduals(fsqr1: Any, fsqz1: Any, fsql1: Any)¶: Preconditioned residuals (VMEC2000 residue.f90: fsqr1/fsqz1/fsql1).

vmex.core.residuals.m1_constrained_to_physical(R_cos: Any, Z_sin: Any, R_sin: Any | None = None, Z_cos: Any | None = None, *, modes: ModeTable, lthreed: bool, lasym: bool, lconm1: bool = True) → tuple[Any, Any, Any | None, Any | None]¶

Undo the m=1 internal constraint on signed spectral coefficients.

VMEC2000: residue.f90 — the evolved (internal) m=1 pairs are rss_int = (rss + zcs)/2 and zcs_int = (rss - zcs)/2 (lconm1, cf. readin.f rbss = (rbs + zbc)/2), so synthesis needs

rss = rss_int + zcs_int, zcs = rss_int - zcs_int

(and rsc/zcc likewise for lasym). This is the mapping the geometry/fields pipeline expects on its inputs. A no-op unless lconm1 and the run is 3D or asymmetric.

Returns (R_cos, Z_sin, R_sin, Z_cos).

vmex.core.residuals.m1_physical_to_constrained(R_cos: Any, Z_sin: Any, R_sin: Any | None = None, Z_cos: Any | None = None, *, modes: ModeTable, lthreed: bool, lasym: bool, lconm1: bool = True) → tuple[Any, Any, Any | None, Any | None]¶: Apply the m=1 internal constraint (inverse of m1_constrained_to_physical()): rss_int = (rss + zcs)/2 etc.

vmex.core.residuals.m1_residue_rotation(force: SpectralForce, *, lconm1: bool = True) → SpectralForce¶

Rotate the m=1 force pairs into the constrained basis (constrain_m1).

VMEC2000: residue.f90 — before getfsq the m=1 rows of the paired blocks are rotated with osqrt2 = 1/sqrt(2):

gcr <- (gcr + gcz)/sqrt(2), gcz <- (gcr - gcz)/sqrt(2)

for the symmetric pair (frss, fzcs) (3D) and the asymmetric pair (frsc, fzcc) (lasym). The conditional zeroing of the constrained component is separate — see zero_m1_z_force() / m1_zero_condition().

vmex.core.residuals.zero_m1_z_force(force: SpectralForce, enabled: Any) → SpectralForce¶

Zero the constrained m=1 Z-force rows (fzcs/fzcc) when enabled.

VMEC2000: residue.f90 (constrain_m1: IF (fsqz < FThreshold .OR. iter2 < 2) gcz = 0). enabled is a traced scalar (bool or 0/1 float, typically m1_zero_condition()), so no retracing occurs when the condition flips.

vmex.core.residuals.m1_zero_condition(*, fsqz_previous: Any, iterations_since_restart: Any) → Any¶

Traced release condition for zero_m1_z_force() (residue.f90).

True (zero the constrained Z force) when iter2 < 2 — counted since the last restart — or the previous fsqz is below M1_FSQZ_RELEASE_THRESHOLD = 1e-6. Pure function of traced values; the caller keeps the values and decides scheduling.

vmex.core.residuals.edge_force_condition(*, fsq_rz_previous: Any, iterations_since_restart: Any, free_boundary: bool) → Any¶

Traced getfsq edge-inclusion rule (residue.f90, SPH041117).

medge = 1 (edge rows join fsqr/fsqz) only for free-boundary runs with iter2 - iter1 < 50 and fsqr + fsqz < 1e-6. Returned as a traced boolean built from the supplied values.

vmex.core.residuals.scalxc_scale_force(force: SpectralForce, *, s: Any) → SpectralForce¶

Apply the odd-m 1/sqrt(s) scaling to every force block.

VMEC2000: funct3d.f — gc = gc * scalxc right after tomnsps, making the scaling part of the definition of the reported residuals. Uses vmex.core.transforms.odd_m_sqrt_s_scaling() (profil3d.f scalxc).

vmex.core.residuals.force_residuals(force: SpectralForce, *, fnorm: Any, fnormL: Any, r1: Any, include_edge: bool = False) → ForceResiduals¶

Invariant residuals fsqr/fsqz/fsql and fedge (getfsq.f).

fsqr = r1*fnorm*sum(gcr**2) and fsqz = r1*fnorm*sum(gcz**2) over surfaces js <= jsmax = ns-1 + medge (include_edge is VMEC’s medge = 1; see edge_force_condition()); fsql = fnormL * sum(gcl**2) over all surfaces. fedge = r1*fnorm times the edge-row R/Z sums (residue.f90) — nonzero only when the projection kept the edge row (tomnsps(include_edge=True)).

The norms come from vmex.core.fields.energies_and_force_norms() (fnorm/fnormL/r1). Pass the force after m1_residue_rotation() / zero_m1_z_force() / scalxc_scale_force(), matching the funct3d.f -> residue.f90 ordering.

vmex.core.residuals.scale_m1_preconditioner_rhs(force: SpectralForce, *, coefficients_R: RadialPreconditionerCoefficients, coefficients_Z: RadialPreconditionerCoefficients, lconm1: bool = True) → SpectralForce¶

m=1 force balance factors before the radial solve (scale_m1).

VMEC2000: residue.f90 (scale_m1_par) — the constrained m=1 rows of the R and Z systems are weighted with the odd-parity preconditioner diagonals so the pair converges at one rate:

fac_R = (ard + brd)/(ard + brd + azd + bzd) (odd column), fac_Z = (azd + bzd)/(ard + brd + azd + bzd),

applied to the frss/fzcs (3D) and frsc/fzcc (lasym) m=1 rows. Following the parity-proven legacy port, the factors scale the right-hand side before apply_radial_preconditioner().

vmex.core.residuals.apply_radial_preconditioner(force: SpectralForce, *, matrices_R: TridiagonalMatrices, matrices_Z: TridiagonalMatrices, jmax: int) → SpectralForce¶

Solve the R and Z radial tridiagonal systems against all force blocks.

VMEC2000: scalfor.f applied to gcr (with the arm/ard/brm/brd/ crd matrices) and gcz (azm/azd/bzm/bzd); each parity family is solved in one batched vmex.core.preconditioner.scalfor() call. Lambda blocks pass through (see apply_lambda_preconditioner()).

vmex.core.residuals.apply_lambda_preconditioner(force: SpectralForce, faclam: Any) → SpectralForce¶: Scale the lambda force blocks by faclam (residue.f90: gcl = faclam*gcl; faclam from vmex.core.preconditioner.lamcal()).

vmex.core.residuals.preconditioned_residuals(force_preconditioned: SpectralForce, *, fnorm1: Any, delta_s: Any) → PreconditionedResiduals¶

Preconditioned residuals fsqr1/fsqz1/fsql1 (residue.f90).

With gc the output of apply_radial_preconditioner() + apply_lambda_preconditioner():

fsqr1 = fnorm1 * sum(gcr**2), fsqz1 = fnorm1 * sum(gcz**2) (all surfaces — the edge row is whatever the projection left there), fsql1 = hs * sum_{js>=2}(gcl**2).

fnorm1 is the spectral-state normalization from vmex.core.fields.preconditioned_force_norm() (bcovar.f); delta_s is the radial spacing hs.

Solver¶

Run setup: radial grids, 1D flux/profile arrays, boundary processing, and the initial interior guess — everything a fixed-boundary solve needs before the first iteration.

VMEC2000 counterparts¶

Sources/Initialization_Cleanup/profil1d.f — radial grids (hs, sqrts, shalf, sm/sp), the 1D flux profiles phips/chips and phipf/chipf, iotas/iotaf (piota), the enclosed-current profile icurv (pcurr + CURTOR scaling) and the mass profile (pmass in internal mu0*Pa units), plus lamscale: radial_grids(), flux_profiles().
Sources/Initialization_Cleanup/magnetic_fluxes.f — torflux/ torflux_deriv from the APHI polynomial and polflux_deriv = piota(tf) * torflux_deriv.
Sources/Input_Output/readin.f — boundary-coefficient processing: the lasym delta rotation, conversion to the internal rbcc/rbss/... blocks, the Jacobian-sign check lflip = (rtest*ztest < 0) with signgs = -1, and the lconm1 m = 1 constraint (rbss = (rbss + zbcs)/2 etc.): boundary_from_input().
Sources/Initialization_Cleanup/init_geometry.f90 — flip_theta (theta -> pi - theta sign factors (-1)**m).
Sources/Initialization_Cleanup/profil3d.f — the interior guess: odd/even scalxc factors and the interpolation of boundary + axis into the volume (rmn(js,m>0) = rmn_bdy * sqrts(js)**m; rmn(js,m=0) = s*rmn_bdy + (1-s)*axis): interior_guess(), run_setup().
Sources/Initialization_Cleanup/guess_axis.f — the axis re-guess grid search used after a first bad-Jacobian start: guess_axis().

Representation conventions (parity-critical)¶

All spectral outputs are in the signed-(m, n) helical packing of vmex.core.fourier.mode_table(), in VMEC internal normalization (divided by mscale(m)*nscale(|n|)) and — for the boundary and the initial state — in the m = 1 constrained basis evolved by the solver (residue.f90 / readin.f lconm1). Geometry synthesis (vmex.core.geometry.real_space_geometry()) consumes the physical basis; geometry_state() performs the conversion (vmex.core.residuals.m1_constrained_to_physical() + the 3D lambda axis closure).

Boundary/axis parsing and the guess_axis grid search are one-time host NumPy code (data-dependent argmax); every produced array is a jnp array and interior_guess() — the state-producing path — is jit-compatible. Ported from the parity-proven legacy implementation (vmex.init_guess, vmex.boundary, vmex.energy, vmex.solvers.fixed_boundary.profiles); equivalence is enforced in tests/test_setup_ab.py.

class vmex.core.setup.RadialGrids(s_full: Any, s_half: Any, sqrts: Any, shalf: Any, sm: Any, sp: Any, hs: Any)¶

VMEC radial full/half meshes (VMEC2000: profil1d.f).

With hs = 1/(ns-1) and 1-based Fortran index i:

s_full(i) = hs*(i-1) — full mesh, s in [0, 1];
s_half(i) = hs*|i-1.5| — half mesh; the axis slot repeats the first interior value (s_half(1) = s_half(2) = hs/2, matching vmex.core.geometry.sqrt_s_half_mesh());
sqrts = sqrt(s_full) with sqrts(ns) = 1 exactly (round-off guard); shalf = sqrt(s_half);
sm(i) = shalf(i)/sqrts(i), sp(i) = shalf(i+1)/sqrts(i) with sm(1) = 0, sp(1) = sm(2) and shalf(ns+1) = 1 — the odd-m half-mesh interpolation weights used by lamcal/precondn;
hs — the uniform radial step (0-d array).

class vmex.core.setup.ProcessedBoundary(R_cos: Any, R_sin: Any, Z_cos: Any, Z_sin: Any, r00: Any, signgs: int, lflip: bool)¶

Boundary coefficients after the readin.f processing chain.

R_cos/R_sin/Z_cos/Z_sin are signed-(m, n) helical coefficients aligned with vmex.core.fourier.mode_table(), in internal normalization (divided by mscale*nscale), theta-flipped when lflip and in the m = 1 constrained basis when lconm1 (readin.f: rbss = (rbss + zbcs)/2 etc.). r00 is the physical RBC(0, 0) (VMEC r00 = rmn_bdy(0,0,rcc)); signgs is always -1 (readin.f); lflip records the theta flip decision rtest*ztest < 0.

class vmex.core.setup.RunSetup(s_full: Any, s_half: Any, sqrts: Any, shalf: Any, sm: Any, sp: Any, hs: Any, scalxc: Any, phips: Any, chips: Any, iotas: Any, icurv: Any, mass: Any, phipf: Any, chipf: Any, iotaf: Any, lamscale: Any, boundary_R_cos: Any, boundary_R_sin: Any, boundary_Z_cos: Any, boundary_Z_sin: Any, raxis_c: Any, raxis_s: Any, zaxis_c: Any, zaxis_s: Any, R_cos: Any, R_sin: Any, Z_cos: Any, Z_sin: Any, lambda_cos: Any, lambda_sin: Any, signgs: int, lflip: bool, lasym: bool, lthreed: bool, lconm1: bool, ncurr: int)¶

Everything a fixed-boundary solve needs before iterating.

Radial grids (profil1d.f; see RadialGrids for definitions): s_full, s_half, sqrts, shalf, sm, sp, hs plus scalxc — the odd-m 1/sqrt(s) factors, shape (ns, mpol) (profil3d.f; equals vmex.core.transforms.odd_m_sqrt_s_scaling()).

1D profiles (profil1d.f, half mesh unless noted; index 0 = axis slot is zeroed exactly as in Fortran):

phips/chips: torflux_edge * torflux_deriv/polflux_deriv(s_half) with torflux_edge = signgs*phiedge/(2*pi) (normalized by torflux(1)); chips (and iotas) are negated when lflip.
iotas: piota(min(torflux(s_half), 1)); iotaf/phipf/chipf: the full-mesh companions (not flipped — profil1d.f quirk).
icurv: Itor * pcurr(tf) with Itor = signgs*mu0*curtor / (2*pi*pcurr(1)) (zero when |pcurr(1)| <= eps*|curtor|). Computed unconditionally; consumed by add_fluxes only when ncurr = 1.
mass: mu0*pres_scale*pmass(tf) * (|phips|*r00)**gamma with the spres_ped clamp (pmass(spres_ped) for s_half > spres_ped).
lamscale = sqrt(hs * sum(phips(2:ns)**2)).

Boundary/axis: boundary_R_cos/... per ProcessedBoundary; raxis_c/raxis_s/zaxis_c/zaxis_s are the physical axis coefficients actually used (input arrays, or the guess_axis() result when the input axis was all-zero and infer_axis_if_missing).

Initial state (profil3d.f): R_cos/R_sin/Z_cos/Z_sin shape (ns, mnmax) — internal-normalized, m = 1-constrained spectral coefficients (the solver’s evolution basis) — and lambda_cos/ lambda_sin (zero; profil1d.f zeroes the lambda block of xc). Use geometry_state() to obtain the physical-basis inputs of vmex.core.geometry.real_space_geometry().

Static metadata: signgs (always -1), lflip, lasym, lthreed, lconm1 and ncurr (0: prescribed iota, 1: prescribed current).

vmex.core.setup.radial_grids(ns: int, *, dtype=<class 'jax.numpy.float64'>) → RadialGrids¶

Build the VMEC radial meshes (VMEC2000: profil1d.f).

s_full(i) = hs*(i-1), s_half(i) = hs*|i-1.5| (1-based i), sqrts/shalf their square roots (sqrts(ns) = 1 exactly), and the sm/sp odd-m interpolation weights (see RadialGrids).

vmex.core.setup.boundary_from_input(inp: VmecInput, *, modes: ModeTable, trig: TrigTables, lconm1: bool = True) → ProcessedBoundary¶

Process INDATA boundary coefficients into the solver representation.

VMEC2000: readin.f — lasym delta rotation, internal-block accumulation, theta flip / signgs = -1 determination, the lconm1 m = 1 constraint, and the mscale*nscale internal normalization applied by profil3d.f (t1 = 1/(mscale(m)*nscale(n))).

Host NumPy (one-time parsing); all outputs are jnp arrays.

vmex.core.setup.flux_profiles(inp: VmecInput, grids: RadialGrids, *, r00: Any, signgs: int = -1, lflip: bool = False) → dict[str, Any]¶

Evaluate the profil1d.f 1D profile arrays on the radial meshes.

VMEC2000: profil1d.f — with torflux_edge = signgs*phiedge/(2*pi) (divided by torflux(1) when nonzero) and tf = min(torflux(s), 1):

half mesh (index 0 zeroed): phips = torflux_edge*torflux_deriv, chips = torflux_edge*piota(tf)*torflux_deriv (polflux_deriv), iotas = piota(tf), icurv = Itor*pcurr(tf) with Itor = signgs*mu0*curtor/(2*pi*pcurr(1)) (0 when |pcurr(1)| <= eps*|curtor|), mass = mu0*pres_scale*pmass * (|phips|*r00)**gamma with the spres_ped clamp;
full mesh: phipf/chipf/iotaf analogously;
lamscale = sqrt(hs*sum(phips(2:)**2));
lflip negates iotas and chips (only — profil1d.f leaves the full-mesh arrays unflipped).

lrfp (RFP mode) is not supported. Returns a dict of jnp arrays keyed phips, chips, iotas, icurv, mass, phipf, chipf, iotaf, lamscale.

vmex.core.setup.interior_guess(*, boundary_R_cos: Any, boundary_R_sin: Any, boundary_Z_cos: Any, boundary_Z_sin: Any, raxis_c: Any, raxis_s: Any, zaxis_c: Any, zaxis_s: Any, modes: ModeTable, trig: TrigTables, s: Any) → tuple[Any, Any, Any, Any, Any, Any]¶

profil3d.f interior interpolation of boundary + axis into the volume.

VMEC2000: profil3d.f (lreset) — starting from xc = 0:

m > 0: coeff(js) = coeff_bdy * sqrts(js)**m (facj);
m = 0: coeff(js) = s(js)*coeff_bdy + (1 - s(js))*axis where the axis coefficients enter internally scaled (rax1*t1); in the signed helical packing the net mapping is R_cos <- raxis_c/nscale, R_sin <- raxis_s/nscale, Z_cos <- zaxis_c/nscale, Z_sin <- zaxis_s/nscale (the two profil3d.f minus signs — the internal zcs/rcs storage and the sin(m*u - n*v) helical basis at m = 0 — cancel);
lambda: zero (profil1d.f zeroes the lambda block of xc).

Boundary coefficients must already be processed by boundary_from_input() (internal-scaled, flipped, m = 1-constrained); the returned state is in the same (evolution) basis. Pure JAX and jit-compatible: modes/trig are static tables, s has a static shape.

Returns (R_cos, R_sin, Z_cos, Z_sin, lambda_cos, lambda_sin), each of shape (ns, mnmax).

vmex.core.setup.guess_axis(geometry: RealSpaceGeometry, *, s: Any, trig: TrigTables, signgs: int = -1, grid_points: int = 61) → tuple[Any, Any, Any, Any]¶

Re-guess the magnetic axis after a bad-Jacobian start (guess_axis.f).

VMEC2000: Sources/Initialization_Cleanup/guess_axis.f — in each zeta plane, scan a limpts x limpts (R, Z) grid over the LCFS bounding box for the axis position maximizing the minimum over theta of the Jacobian proxy

tau = signgs * (ru12*(zs + z_axis') - zu12*(rs + r_axis'))

built from the boundary surface (js = ns) and the mid surface (js = ns12 = (ns+1)/2) with rs = (r1b - r12)/ds + r_axis; the symmetric case mirrors theta in (pi, 2*pi) from the reduced grid and scans only half the zeta planes, and up-down-symmetric planes fix z = 0. The winning per-plane positions are Fourier-projected with the nscale-normalized zeta tables (halving n = 0 and n = nzeta/2).

geometry supplies the totzsps channels (r1/z1 even/odd with the internal odd representation, and the theta derivatives for ru0/zu0) — exactly the inputs of the Fortran routine. Host NumPy code (data-dependent grid search; runs once after a failed start), ported verbatim from the parity-proven vmex.init_guess._recompute_axis_from_state_vmec.

Returns physical axis coefficient arrays (raxis_c, raxis_s, zaxis_c, zaxis_s), each of length ntor + 1.

vmex.core.setup.geometry_state(setup: RunSetup, *, modes: ModeTable) → dict[str, Any]¶

Convert the evolution-basis state into geometry-synthesis inputs.

Applies vmex.core.residuals.m1_constrained_to_physical() (residue.f90 — undo the internal m = 1 constraint) and the 3D lambda axis closure (totzsp_mod.f); the result is the keyword dict expected by vmex.core.geometry.real_space_geometry().

vmex.core.setup.run_setup(inp: VmecInput, resolution: Resolution, *, lconm1: bool = True, infer_axis_if_missing: bool = True) → RunSetup¶

Build the complete pre-iteration setup for one radial resolution.

VMEC2000: the readin.f boundary processing + profil1d.f + profil3d.f initialization sequence of runvmec.f/eqsolve.f.

Parameters:

inp – Parsed input (vmex.core.input.VmecInput); its fields are consumed as concrete host numbers (setup runs once per solve).
resolution – Static resolution; resolution.ns selects the multigrid stage.
lconm1 – Apply the m = 1 constraint conversion (VMEC2000 default T).
infer_axis_if_missing – When every input axis coefficient is zero, run guess_axis() on the zero-axis interior guess and re-blend (the legacy driver default; VMEC2000 itself would start from the zero axis and only call guess_axis after the first bad-Jacobian restart).

Return type:

RunSetup (see its docstring for the field-by-field contract).

1D (radial) preconditioner for the VMEC R/Z/lambda force equations.

VMEC2000 counterparts¶

Sources/General/precondn.f — precondn(): radial tridiagonal matrix elements for the R and Z force equations from ptau-type flux-surface integrals. VMEC calls it twice per refresh: once with the Z geometry derivatives (producing arm/ard/brm/brd/crd for the R force) and once with the R geometry derivatives (producing azm/azd/bzm/bzd for the Z force).
Sources/General/lamcal.f90 — lamcal(): the diagonal lambda-force preconditioner faclam.
Sources/General/scalfor.f — scalfor_matrices() (assembly of the tridiagonal system ax/bx/dx with the m**2 and (n*nfp)**2 mode weights, edge_pedestal and the ZC(0,0)(ns) stabilization) and scalfor() (application: solve the system in place of the force).
Sources/General/tridslv (serial_tridslv in scalfor.f) — tridiagonal_solve(): Thomas algorithm, vectorized over all (m,n) columns and stacked right-hand-side fields simultaneously.

Equations (VMEC2000 notation)¶

precondn builds, for the linearized radial force operator of the energy functional, the flux-surface integrals over the half mesh (js = 2..ns):

ptau = pfactor * r12**2 * bsq * wint / gsqrt          (pfactor = -4)

ax = < ptau * (xu12/hs ± t23)**2 >    (poloidal-derivative couplings)
bx = < ptau * t1 * t2 >               (radial-derivative couplings)
cx = < 0.25 * pfactor * (bsupv**2) * gsqrt >   (toroidal couplings)

which are then accumulated onto the full mesh as off-diagonal (axm/bxm, between surfaces) and diagonal (axd/bxd, per surface) elements, each with an even-m column and an odd-m column (odd-m carries the sm/sp internal sqrt(s) scalings from profil1d).

scalfor forms, per mode (m,n), the radial tridiagonal system:

ax(js) = -(axm(js+1) + bxm(js+1) * m**2)          couples X(js+1)
bx(js) = -(axm(js)   + bxm(js)   * m**2)          couples X(js-1)
dx(js) = -(axd(js) + bxd(js)*m**2 + cx(js)*(n*nfp)**2)

with the even/odd column selected by mod(m,2), and solves [bx, dx, ax] . X = force by the Thomas algorithm (tridslv).

lamcal builds the diagonal lambda scaling:

blam = < guu / gsqrt >,  clam = < gvv / gsqrt >,  dlam = < guv / gsqrt >
faclam = p_factor * sqrt(s)**power /
         (blam*(n*nfp)**2 + sign(dlam,blam)*2*m*(n*nfp) + clam*m**2)

with power = min(m**2/16**2, 8) (the sqrt(s) damping only bites for m > 16) and a special (m,n)=(0,0) slot that preconditions the chip/iota channel (VMEC stores chip-like data in lmns(0,0)).

Index conventions (parity-critical)¶

Radial arrays are 0-based here: full-mesh index js = 0..ns-1 maps to VMEC’s 1..ns; half-mesh arrays are stored without VMEC’s dummy axis row, so row j of a half-mesh input is VMEC’s js = j+2 half point (between full points j and j+1). Exception: lamcal() takes the half-mesh metrics with the dummy axis row (ns rows), exactly as VMEC’s guu/gvv/guv(1:ns) storage, and internally copies row 1 into row 0 (blam(1) = blam(2) in lamcal.f90).
jmax: number of evolved radial rows. Fixed boundary: jmax = ns-1 (VMEC jmax = ns1 when ivac < 1); the edge row keeps the raw force. Free boundary: jmax = ns and the edge_pedestal/ZC(0,0) stabilization of scalfor.f activates (it only applies when jmax >= ns).
jmin (start row of the tridiagonal solve, VMEC jmin2): 0 for m = 0 and 1 for m >= 1; row 0 of the m >= 1 solution is zeroed. For m = 1 the sub-diagonal coupling of row 1 to the axis row is folded into the diagonal (dx(2,m=1) += bx(2,m=1), reflecting VMEC’s X(1,m=1) = X(2,m=1) axis convention).
lambda start index (VMEC jlam(m) = 2): row 0 of faclam is zero for every mode except (0,0).

All functions are pure jax.numpy (jit-friendly, no host round-trips); resolution parameters (ns/mpol/ntor/nfp/...) are Python ints/bools and must be static under jax.jit. The numerics are ported verbatim from the parity-proven legacy module vmex.preconditioner_1d_jax; equivalence is enforced in tests/test_preconditioner_ab.py.

class vmex.core.preconditioner.RadialPreconditionerCoefficients(axm: Array, axd: Array, bxm: Array, bxd: Array, cx: Array)¶

Radial matrix elements from precondn() (VMEC2000 precondn.f).

Columns index the m-parity: [:, 0] even m, [:, 1] odd m (odd-m carries the internal sqrt(s) scalings sm/sp). At the R-force call site these are VMEC’s arm/ard/brm/brd/crd; at the Z-force call site azm/azd/bzm/bzd (cx is shared).

axm: Any¶: (ns-1, 2) off-diagonal poloidal coupling (half-mesh rows).

axd: Any¶: (ns, 2) diagonal poloidal coupling (full-mesh accumulated).

bxm: Any¶: (ns-1, 2) off-diagonal radial coupling (half-mesh rows).

bxd: Any¶: (ns, 2) diagonal radial coupling (full-mesh accumulated).

cx: Any¶: (ns,) toroidal coupling (full-mesh accumulated).

class vmex.core.preconditioner.TridiagonalMatrices(ax: Array, bx: Array, dx: Array)¶

Assembled per-mode radial tridiagonal system (VMEC2000 scalfor.f).

ax[js] couples row js to js+1 (superdiagonal), bx[js] to js-1 (subdiagonal), dx is the diagonal. Shapes (jmax, mpol, ntor+1).

ax: Any¶: Alias for field number 0

bx: Any¶: Alias for field number 1

dx: Any¶: Alias for field number 2

vmex.core.preconditioner.angular_integration_weights(*, ntheta: int, nzeta: int, lasym: bool) → ndarray¶

Poloidal integration weights for the flux-surface averages, (ntheta_eff,).

VMEC2000: wint from profil3d.f (constant in zeta), i.e. dnorm3 with half weights at the theta = 0, pi endpoints of the reduced grid for stellarator-symmetric runs; uniform 1/(nzeta*ntheta1) on the full grid when lasym. Matches TrigTables.wint[:, 0] of vmex.core.fourier. Static (NumPy) by construction.

vmex.core.preconditioner.precondn(*, dxds_half: Any, dxdu_half: Any, dxdu_even_full: Any, dxdu_odd_full: Any, x_odd_full: Any, r12_half: Any, bsq_half: Any, bsupv_half: Any, sqrt_g_half: Any, angular_weight: Any, delta_s: Any, ns: int) → RadialPreconditionerCoefficients¶

Radial tridiagonal matrix elements for one force family (precondn.f).

For the R force pass the Z geometry derivatives (VMEC: precondn(..., zs, zu12, zu(even), zu(odd), z1(odd), ...) producing arm/ard/brm/brd/crd); for the Z force pass the R derivatives (producing azm/azd/bzm/bzd).

Parameters:

dxds_half – (ns-1, ntheta, nzeta) radial derivative of Z (or R) on the half mesh, no axis dummy row (VMEC zs/rs).
dxdu_half – (ns-1, ntheta, nzeta) poloidal derivative on the half mesh (VMEC zu12/ru12).
dxdu_even_full – (ns, ntheta, nzeta) even-m / odd-m (internal) poloidal-derivative channels on the full mesh (VMEC zu(:,0)/zu(:,1)).
dxdu_odd_full – (ns, ntheta, nzeta) even-m / odd-m (internal) poloidal-derivative channels on the full mesh (VMEC zu(:,0)/zu(:,1)).
x_odd_full – (ns, ntheta, nzeta) odd-m channel of Z (or R) itself (VMEC z1(:,1)).
r12_half – (ns-1, ntheta, nzeta) half-mesh R, total pressure |B|**2/2 + p (VMEC bsq), contravariant B^v and Jacobian gsqrt, no axis dummy row.
bsq_half – (ns-1, ntheta, nzeta) half-mesh R, total pressure |B|**2/2 + p (VMEC bsq), contravariant B^v and Jacobian gsqrt, no axis dummy row.
bsupv_half – (ns-1, ntheta, nzeta) half-mesh R, total pressure |B|**2/2 + p (VMEC bsq), contravariant B^v and Jacobian gsqrt, no axis dummy row.
sqrt_g_half – (ns-1, ntheta, nzeta) half-mesh R, total pressure |B|**2/2 + p (VMEC bsq), contravariant B^v and Jacobian gsqrt, no axis dummy row.
angular_weight – (ntheta,) or (ntheta, nzeta) integration weights (angular_integration_weights(); VMEC wint).
delta_s – Radial mesh spacing hs = 1/(ns-1) (VMEC hs; ohs = 1/hs).
ns – Number of full-mesh surfaces (static).

vmex.core.preconditioner.lamcal(*, guu_half: Any, guv_half: Any, gvv_half: Any, sqrt_g_half: Any, lamscale: Any, angular_weight: Any, mpol: int, ntor: int, nfp: int, lthreed: bool, damping_factor: float = 2.0, r0scale: float = 1.0) → Any¶

Diagonal lambda-force preconditioner faclam (VMEC2000 lamcal.f90).

Parameters:

guu_half – (ns, ntheta, nzeta) half-mesh metric elements and Jacobian including the dummy axis row (VMEC storage); the axis row is replaced internally (blam(1) = blam(2)). guv is ignored unless lthreed.
guv_half – (ns, ntheta, nzeta) half-mesh metric elements and Jacobian including the dummy axis row (VMEC storage); the axis row is replaced internally (blam(1) = blam(2)). guv is ignored unless lthreed.
gvv_half – (ns, ntheta, nzeta) half-mesh metric elements and Jacobian including the dummy axis row (VMEC storage); the axis row is replaced internally (blam(1) = blam(2)). guv is ignored unless lthreed.
sqrt_g_half – (ns, ntheta, nzeta) half-mesh metric elements and Jacobian including the dummy axis row (VMEC storage); the axis row is replaced internally (blam(1) = blam(2)). guv is ignored unless lthreed.
lamscale – VMEC lambda scale factor (bcovar: sqrt(sum phip**2 * wint)).
angular_weight – As in precondn() (VMEC wint).
mpol – Mode-space resolution (static).
ntor – Mode-space resolution (static).
nfp – Mode-space resolution (static).
lthreed – Mode-space resolution (static).
damping_factor – VMEC uses 2.0 and mscale(0)*nscale(0) = 1.
r0scale – VMEC uses 2.0 and mscale(0)*nscale(0) = 1.

Returns:

``(ns, mpol, ntor+1)`` array ``faclam`` (the multiplicative preconditioner)
of the lambda force, with
- sqrt(s)**min(m**2/256, 8) damping (only active for m > 16),
- the (0,0) slot set to damping/(4*r0scale**2) / blam — VMEC’s – chip/iota channel (the lamscale**2 in p_factor cancels there),
- the axis row zeroed except (0,0) (VMEC jlam(m) = 2).

vmex.core.preconditioner.scalfor_matrices(coefficients: RadialPreconditionerCoefficients, *, delta_s: Any, mpol: int, ntor: int, nfp: int, ns: int, jmax: int | None = None, stabilize_edge_zc00: bool = False) → TridiagonalMatrices¶

Assemble the per-mode tridiagonal system (matrix half of scalfor.f).

Parameters:

coefficients – Output of precondn() — VMEC’s (arm, ard, brm, brd, crd) for the R force or (azm, azd, bzm, bzd, crd) for the Z force.
delta_s – Radial spacing hs (only used by the edge stabilization factor).
mpol – Static resolution parameters.
ntor – Static resolution parameters.
nfp – Static resolution parameters.
ns – Static resolution parameters.
jmax – Number of evolved rows; defaults to ns-1 (fixed boundary, VMEC jmax = ns1). Pass ns for free-boundary-style solves — only then do the EDGE_PEDESTAL and ZC(0,0) edge terms activate.
stabilize_edge_zc00 – Apply the ZC(0,0)(ns) fac = 0.25 stabilization (mult_fac = min(fac, fac*hs*15)); VMEC applies it to the Z force system only (iflag = 1 call site, .not. lfreeb).

Notes

Assembly per mode (VMEC 1-based js, even/odd column mod(m,2)):

ax(js) = -(axm(js+1,mp) + bxm(js+1,mp)*m**2)
bx(js) = -(axm(js,mp)   + bxm(js,mp)*m**2)
dx(js) = -(axd(js,mp) + bxd(js,mp)*m**2 + cx(js)*(n*nfp)**2)

with the axis row zeroed for m >= 1, the m=1 axis fold dx(2,m=1) += bx(2,m=1), and (edge rows only, jmax >= ns) the diagonal stiffened by 1+EDGE_PEDESTAL (m <= 1) or 1+2*EDGE_PEDESTAL (m >= 2).

vmex.core.preconditioner.scalfor(force: Any, matrices: TridiagonalMatrices, *, jmax: int, tridiagonal_method: str = 'auto') → Any¶

Apply the assembled preconditioner to a spectral force (scalfor.f).

Solves the radial tridiagonal systems of matrices against force respecting the VMEC start-index conventions (jmin2): the m = 0 block is solved on rows 0..jmax-1, the m >= 1 blocks on rows 1..jmax-1 with the axis row of the solution set to zero. Rows >= jmax (the fixed edge in fixed-boundary mode) keep the input force.

Parameters:

force – (ns, mpol, ntor+1) spectral force (VMEC gcr/gcz), or (ns, mpol, ntor+1, nfields) to precondition several parity channels (e.g. frcc/frss or fzsc/fzcs) in one batched solve.
matrices – Output of scalfor_matrices() for the matching force family.
jmax – Same jmax used to assemble matrices (static).
tridiagonal_method – Forwarded to tridiagonal_solve() ("auto"/"thomas"/ "lax"; static). The default picks per lowering platform: the bit-parity Thomas scan on CPU, the fused cuSPARSE-backed kernel on accelerators.

vmex.core.preconditioner.tridiagonal_solve(superdiagonal: Any, diagonal: Any, subdiagonal: Any, rhs: Any, *, method: str = 'auto') → Any¶

Radial tridiagonal solve vectorized over all (m,n) columns and fields.

VMEC2000: serial_tridslv (in scalfor.f): solve, along the leading (radial) axis, sub[j]*x[j-1] + diag[j]*x[j] + super[j]*x[j+1] = rhs[j] for every trailing-column element simultaneously. sub[0] and super[-1] are ignored. rhs may carry extra trailing axes (stacked force fields) beyond the shape of diagonal; they are solved in one pass.

Thin arg-order adapter over solvax.tridiagonal_solve() (roadmap R18b shared-solver consolidation). SOLVAX uses the (lower, diag, upper) band convention (lower = sub, upper = super), so the sub-/super- diagonals are swapped here to preserve vmex’s (super, diag, sub) signature and the two scalfor call sites. The numerics are unchanged: SOLVAX’s Thomas backend is the verbatim port of the legacy parity-proven sweep (same eps = 1e-12), so method still selects it bit-for-bit —

"thomas": two lax.scan Thomas sweeps (jit-friendly, no host round-trips); the CPU path, bitwise identical to the pinned A/B tests.
"lax": XLA’s fused batched solver jax.lax.linalg.tridiagonal_solve (cuSPARSE gtsv2 on CUDA, LAPACK gtsv on CPU); the accelerator fast path (numerically equivalent, not bit-identical to Thomas).
"auto" (default): Thomas when lowering for CPU (bit parity), the fused solver elsewhere, chosen with jax.lax.platform_dependent. Systems with fewer than 3 radial rows always use Thomas.

2D block preconditioner: a matrix-free Newton step via jax.jvp + SOLVAX.

VMEC2000 counterpart¶

Sources/Hessian/precon2d.f — VMEC2000’s optional 2D (ictrl_prec2d) preconditioner. It builds the full block-tridiagonal-in-radius Hessian of the 1D-preconditioned force map gc = g(xc) by finite-difference “jogs” of every (mode, type) column (compute_blocks), LU-factors it with BCYCLIC (blk3d_factor), and in residue replaces the force by the block-solve gc <- -H^{-1} gc (block_precond, where H = d gc / d xc). This is a Newton step on the already-1D-preconditioned residual; it converges stiff cases (high beta, high aspect, high mode number) in far fewer iterations than the 1D radial preconditioner alone. Activation (evolve.f): finest grid only, once fsqr + fsqz + fsql < prec2d_threshold and iter2 >= 10.

Modern (JAX) approach¶

The forces are already traceable, so we get exact Hessian-vector products for free from jax.jvp — no finite-difference jogs, no assembled blocks, no BCYCLIC factorization. Let g be the 1D-preconditioned force map state -> gc (frozen ns4 cache — the 1D operator is a fixed linear preconditioner during the Newton solve). The Newton direction that drives g -> 0 is

delta = -J^{-1} g, J = d g / d state (block-tridiagonal in radius),

obtained by solving J delta = -g with matrix-free GMRES (solvax.gmres()) whose matvec is v -> jvp(g, state, v). Because the 1D preconditioner is baked into g (and hence into both J and the right-hand side -g), the M_1D^{-1} factor cancels exactly out of the solve: delta is the same full Newton step -(dF/dstate)^{-1} F on the raw force F regardless of the 1D operator — the 1D preconditioner only accelerates GMRES (it makes J close to the identity near equilibrium, where M_1D approximates dF/dstate). This matches block_precond’s sign: gc_out = -H^{-1} gc with H = d gc / d xc (precon2d.f factorization check block_dsave * x = -gc_save).

The module is deliberately VMEC-agnostic: newton_direction() takes any pytree-valued force map and a linearization point, so it is unit-testable against a dense reference (tests/test_preconditioner_2d.py) and could equally drive the block-tridiagonal solvax.block_thomas_truncated() route if the blocks were assembled explicitly (they are not, by design — the matrix-free HVP keeps peak memory at one force graph).

Wiring lives in vmex.core.solver (_make_body): when precon_type != "NONE" the traced iteration replaces the 1D force direction by newton_direction() under a lax.cond gated on the activation predicate, so the default 1D-only path is untouched.

class vmex.core.preconditioner_2d.Prec2DConfig(threshold: float, start_iteration: int = 10, step: float = 1.0, gmres_restart: int = 40, gmres_max_restarts: int = 3, gmres_rtol: float = 0.01, gmres_atol: float = 0.0, finest: bool = True)¶

Static configuration of the 2D block preconditioner (hashable meta).

threshold¶

Activate the Newton step once fsqr + fsqz + fsql < threshold on the finest grid (VMEC2000 prec2d_threshold / evolve.f). The default input value 1e-30 means “never” — the caller supplies a finite threshold to switch it on.

Type:: float

start_iteration¶

Earliest iteration at which the Newton step may activate (VMEC2000 skips the first 10 iterations: IF (iter2 < 10) ictrl_prec2d = 0).

Type:: int

step¶

Damping applied to the (full) Newton step, state += step * delta. 1.0 is the undamped Newton update; VMEC2000 instead folds the step into its second-order Richardson stepper with time_step = 0.5.

Type:: float

gmres_restart¶

Arnoldi cycle size m of the inner GMRES (clamped to the problem size for tiny cases).

Type:: int

gmres_max_restarts¶

Maximum number of GMRES restart cycles (outer iterations).

Type:: int

gmres_rtol, gmres_atol: Inner GMRES tolerances on ||b - J delta||. A loose rtol gives an inexact Newton step (cheaper, still super-linear far from the asymptotic regime); tighten for a near-exact step.

finest¶

Whether this runtime is the finest multigrid stage (VMEC2000 ns == ns_maxval). Single-grid solves are always finest; only the last multigrid stage sets this True.

Type:: bool

vmex.core.preconditioner_2d.flat_operator(force_map: Callable[[Any], Any], x0: Any) → tuple[Callable[[Any], Any], Callable[[Any], Any], int]¶

Flat matrix-free Jacobian operator of force_map linearized at x0.

Returns (matvec, unravel, n) where matvec(v_flat) = J @ v_flat with J = d force_map / d x evaluated at x0 (one jax.jvp — an exact Hessian-vector product, no finite differences), unravel reconstructs the pytree from a flat vector, and n is the flattened dimension. The output pytree of force_map must share the structure of x0 (a square Jacobian), as it does for the VMEC force map (state channels -> force channels).

vmex.core.preconditioner_2d.newton_direction(force_map: Callable[[Any], Any], x0: Any, rhs: Any, cfg: Prec2DConfig, *, guess: Any | None = None)¶

Matrix-free Newton direction: solve J delta = rhs for delta.

J = d force_map / d x at x0 (exact HVP via flat_operator()), solved with restarted flexible GMRES (solvax.gmres()). For the VMEC Newton step pass force_map = g (the 1D-preconditioned force), x0 = state and rhs = -g(state) so that delta = -J^{-1} g is the descent-consistent Newton update used in place of the raw force.

Parameters:

force_map – As in flat_operator().
x0 – As in flat_operator().
rhs – Right-hand side pytree, same structure as x0.
cfg – Prec2DConfig supplying the GMRES parameters.
guess – Optional initial-guess pytree for delta (defaults to zeros).

Returns:

``(delta, solution)`` (the Newton-direction pytree and the raw)
solvax.krylov.KrylovSolution (residual_norm, iterations,
converged) for diagnostics.

Damped second-order Richardson time stepping and restart control.

Advances the spectral state xc with the momentum (“conjugate gradient without line searches”, after P. Garabedian) scheme of VMEC2000, including the ndamp-window adaptive damping, the Jacobian-reset back-off, and the residual-growth back-off.

VMEC2000 counterparts: Sources/TimeStep/evolve.f (the dtau damping recurrence and momentum update, and TimeStepControl), and Sources/TimeStep/restart.f (restart_iter state restore and time-step rescaling). The parity-critical constants are recorded in Appendix D and asserted here:

damping window NDAMP = 10; per-step decrement capped at 0.15
velocity update v' = (1-dtau)/(1+dtau) * v + delt * F; x' = x + delt*v'
Jacobian reset (irst=2): restore state, zero velocity, delt *= 0.90
residual-growth back-off (irst=3): growth > 1e4 x best after > 10 steps since the last reset; restore state, zero velocity, delt /= 1.03

All functions are pure and jit-compatible; conditions are traced values (no Python branching on array data), so a single compiled step serves both the differentiable and the CLI execution lanes (§5.3).

vmex.core.step.DAMPING_CAP = 0.15¶: Cap on the per-step damping decrement (VMEC2000 cp15 in evolve.f).

vmex.core.step.JACOBIAN_RESET_FACTOR = 0.9¶: Time-step factor on a Jacobian reset (restart.f, irst=2).

vmex.core.step.GROWTH_BACKOFF_DIVISOR = 1.03¶: Time-step divisor on residual growth (restart.f, irst=3).

vmex.core.step.GROWTH_LIMIT = 10000.0¶: Residual-growth ratio that triggers irst=3 (evolve.f TimeStepControl).

vmex.core.step.GROWTH_MIN_ITERATIONS = 10¶: Minimum iterations since the last reset before irst=3 can trigger.

class vmex.core.step.StepControl(time_step: Array, inv_tau: Array, fsq_total_prev: Array, residual_best_precond: Array, residual_best_raw: Array, iter_last_reset: Array, jacobian_resets: Array)¶

Traced scalar state of the time-step controller.

Attributes mirror VMEC2000 module variables: time_step (delt), inv_tau (otau damping history), fsq_total_prev (fsq1), residual_best_precond / residual_best_raw (res0/res1), iter_last_reset (iter1), jacobian_resets (ijacob).

vmex.core.step.damping_coefficients(control: StepControl, iteration: Array, fsq_total: Array) → tuple[Array, Array, StepControl]¶

Advance the ndamp damping window and return (b1, fac) (evolve.f).

fsq_total is the current preconditioned residual sum fsqr1 + fsqz1 + fsql1. On the first step after a (re)start (iteration == iter_last_reset) the window resets to the cap.

vmex.core.step.momentum_update(xc, xcdot, force, b1: Array, fac: Array, time_step: Array)¶

One momentum step: v’ = fac*(b1*v + delt*F); x’ = x + delt*v’ (evolve.f).

xc/xcdot/force may be any matching pytree of arrays.

vmex.core.step.restart_decision(control: StepControl, iteration: Array, fsq_raw_total: Array, fsq_precond_total: Array, jacobian_sign_changed: Array) → tuple[Array, StepControl]¶

Classify the step: STEP_OK, RESTART_JACOBIAN, or RESTART_GROWTH.

Ports evolve.f TimeStepControl: track the best raw/preconditioned residuals seen since the last reset; flag irst=3 when either grows by more than GROWTH_LIMIT x best after GROWTH_MIN_ITERATIONS steps. A Jacobian sign change always wins (irst=2).

vmex.core.step.apply_restart(xc, xcdot, xc_saved, control: StepControl, kind: Array, iteration: Array)¶

Apply restart.f: restore saved state, zero velocity, rescale delt.

For STEP_OK this is the identity on (xc, xcdot) and records the current state as the new save point. Returns (xc, xcdot, xc_saved, control).

Single-grid fixed-boundary solve loop: funct3d evaluation + eqsolve iteration.

Wires the ported core modules (geometry, fields, forces, residuals, preconditioner, step) into one force evaluation (evaluate_forces()) and the damped Richardson iteration (solve()) with VMEC2000 restart/escalation semantics.

VMEC2000 counterparts¶

Sources/General/funct3d.f — one force evaluation: odd-m internal synthesis (totzsps), constraint baselines rcon0/zcon0, Jacobian (early exit on sign change), bcovar (with the ns4 = 25-iteration refresh of the preconditioner / force norms / tcon), constraint force (alias), MHD forces, tomnsps, gc = gc*scalxc and residue: evaluate_forces().
Sources/TimeStep/evolve.f — damping window + momentum update (via vmex.core.step) and TimeStepControl (store/restore of the best state, residual-growth back-off irst = 3): the iteration body.
Sources/TimeStep/eqsolve.f — the outer force-iteration loop: axis re-guess after a first bad Jacobian (ijacob == 0), time-step resets at ijacob = 25/50, abort at ijacob >= 75 (jac75_flag), convergence when fsqr, fsqz, fsql <= ftolv simultaneously: solve().
Sources/TimeStep/restart.f — state restore semantics (zero velocity, delt rescaling) via vmex.core.step.

Execution lanes (§5.3)¶

Both lanes share one traced single-iteration body:

mode="cli": Python while over a jitted lax.scan block of block_size = 10 iterations, with host residual checks and VMEC2000 screen printing (printout.f cadence) between blocks;
mode="jit": a single lax.while_loop over the same body (fully traced iteration; the setup and error raising remain host code).

The per-iteration (fsqr, fsqz, fsql, fsqr1, fsqz1, fsql1, ...) trajectory is recorded in a preallocated buffer carried through the loop, so the two lanes produce identical histories.

Structural executable reuse (2026-07-09, Phase 2 item (1))¶

SolverRuntime is a registered pytree passed as an argument to the module-level jitted lanes _while_lane()/_block_lane() (previously per-runtime closures cached by object identity). Array-valued run data (RunSetup, rcon0/zcon0) are pytree data; the hashable configuration (Resolution, gamma/tcon0/ftol/max_iterations/time_step0/nstep/jmax) is pytree meta; the NumPy mode/trig/weight/gather tables — consumed with np.*/fancy indexing at trace time — are derived from the meta resolution through the cached _static_tables() and never enter the pytree. Consequence: two different runtimes with equal structure (e.g. different boundary values at one resolution, hot restarts, multigrid re-runs) share one XLA executable per lane. solve() additionally accepts initial_state (hot restart), and the stage machinery is factored into _solve_stage()/ _finalize() for reuse by vmex.core.multigrid.solve_multigrid().

vmex.core.solver.NS4 = 25¶

preconditioner / force-norm / tcon refresh cadence.

Type:: bcovar.f ns4

class vmex.core.solver.SpectralState(R_cos: Any, R_sin: Any, Z_cos: Any, Z_sin: Any, L_cos: Any, L_sin: Any)¶

Evolved spectral state in the signed-(m, n) internal packing.

VMEC2000 xc: internal normalization (mscale*nscale divided out), m = 1-constrained basis (residue.f90 lconm1), odd-m WITHOUT the scalxc factor (totzsps applies it during synthesis). For stellarator-symmetric runs R_sin/Z_cos/L_cos are identically zero. All arrays have shape (ns, mnmax).

class vmex.core.solver.PreconditionerCache(tcon: Any, fnorm: Any, fnormL: Any, fnorm1: Any, coefficients_R: RadialPreconditionerCoefficients, coefficients_Z: RadialPreconditionerCoefficients, matrices_R: TridiagonalMatrices, matrices_Z: TridiagonalMatrices, faclam: Any)¶

Quantities refreshed on the ns4 = 25 cadence (VMEC2000 bcovar.f).

tcon (constraint scaling), the residual norms fnorm/fnormL/fnorm1, the precondn() coefficients and assembled scalfor_matrices() for the R and Z force families, and the lamcal diagonal faclam.

class vmex.core.solver.FunctDiagnostics(preconditioned: PreconditionedResiduals, wb: Any, wp: Any, r00: Any, z00: Any, jacobian_sign_changed: Any, cache: PreconditionerCache)¶

Per-evaluation diagnostics returned by evaluate_forces().

preconditioned are the residue.f90 fsqr1/fsqz1/fsql1; wb/wp the MHD energies (bcovar.f); r00/z00 the axis position at theta = zeta = 0 (funct3d.f); jacobian_sign_changed the jacobian.f irst = 2 flag; cache the (possibly refreshed) PreconditionerCache to carry into the next iteration.

class vmex.core.solver.SolverRuntime(resolution: Resolution, setup: RunSetup, rcon0: Any, zcon0: Any, gamma: float, tcon0: float, ftol: float, max_iterations: int, time_step0: float, nstep: int, jmax: int, lfreeb: bool = False, bsqvac_edge: Any | None = None, presf_ns_scale: Any | None = None, prec2d: Any = None)¶

Per-solve context, registered as a JAX pytree.

Passed as an argument to the module-level jitted lanes (_while_lane() / _block_lane()):

data fields (traced): the RunSetup arrays (profiles, radial grids, boundary, initial state — everything that changes with boundary values) and the fixed-boundary constraint baselines rcon0/zcon0 (funct3d.f — constant per run because the edge spectral row never evolves, but boundary-value dependent);
meta fields (static, hashable): the Resolution plus the scalar configuration (gamma is consumed concretely by fields.magnetic_fields; max_iterations sizes the trajectory buffer; the rest are loop-control constants).

The NumPy mode/trig/weight/gather tables are derived from the meta resolution via the cached _static_tables() (exposed as properties), so they never enter the pytree: they are used with np.*/fancy indexing at trace time and must stay concrete. Two runtimes with equal structure (same resolution + scalars, same array shapes) therefore share one XLA executable per lane.

class vmex.core.solver.SolveResult(converged: bool, iterations: int, ier_flag: int, fsqr: float, fsqz: float, fsql: float, wb: float, wp: float, wmhd: float, r00: float, time_step: float, jacobian_resets: int, state: SpectralState, xm: ndarray, xn: ndarray, rmnc: ndarray, zmns: ndarray, rmns: ndarray | None, zmnc: ndarray | None, iotaf: ndarray, fsq_history: ndarray)¶

Converged fixed-boundary solve output.

rmnc/zmns (+ rmns/zmnc when lasym) are physical (wout convention) spectral coefficients on the full mesh, mode-ordered like the wout xm/xn arrays; iotaf follows add_fluxes.f90 for ncurr = 1. fsq_history has one row per iteration: (fsqr, fsqz, fsql, fsqr1, fsqz1, fsql1). wmhd is the printed WMHD = (wb + wp/(gamma-1)) * (2 pi)^2.

class vmex.core.solver.Prec2DConfig(threshold: float, start_iteration: int = 10, step: float = 1.0, gmres_restart: int = 40, gmres_max_restarts: int = 3, gmres_rtol: float = 0.01, gmres_atol: float = 0.0, finest: bool = True)¶

Static configuration of the 2D block preconditioner (hashable meta).

threshold¶

Type:: float

start_iteration¶

Earliest iteration at which the Newton step may activate (VMEC2000 skips the first 10 iterations: IF (iter2 < 10) ictrl_prec2d = 0).

Type:: int

step¶

Type:: float

gmres_restart¶

Arnoldi cycle size m of the inner GMRES (clamped to the problem size for tiny cases).

Type:: int

gmres_max_restarts¶

Maximum number of GMRES restart cycles (outer iterations).

Type:: int

gmres_rtol, gmres_atol: Inner GMRES tolerances on ||b - J delta||. A loose rtol gives an inexact Newton step (cheaper, still super-linear far from the asymptotic regime); tighten for a near-exact step.

finest¶

Whether this runtime is the finest multigrid stage (VMEC2000 ns == ns_maxval). Single-grid solves are always finest; only the last multigrid stage sets this True.

Type:: bool

vmex.core.solver.resolution_from_input(inp: VmecInput, *, ns: int | None = None) → Resolution¶

Resolve the internal grid sizes from an input (read_indata.f).

ntheta <= 0 -> 2*mpol + 6; nzeta = 1 for axisymmetric inputs with nzeta = 0, else nzeta <= 0 -> 2*ntor + 4. ns defaults to the first ns_array stage (single-grid solve).

Build the static solver context from an input file or a RunSetup.

Defaults come from the input (delt, tcon0, gamma, nstep, ftol_array(1), niter_array(1)); explicit keywords override. The fixed-boundary constraint baselines rcon0/zcon0 = s * rcon(ns) (funct3d.f) are computed once here from the initial state — the edge spectral row never evolves in fixed-boundary mode, so they are constants of the run.

precon_type/prec2d_threshold (or an explicit prec2d Prec2DConfig) switch on the optional 2D block preconditioner (precon2d.f); None/"NONE" (the default) leaves the 1D-only path byte-identical.

vmex.core.solver.evaluate_forces(state: SpectralState, runtime: SolverRuntime, *, cache: PreconditionerCache | None = None, iteration: int = 1, iter_last_reset: int = 1, fsqz_previous: float = 1.0) → tuple[SpectralState, ForceResiduals, FunctDiagnostics]¶

One funct3d pass: preconditioned force gc, residuals, diagnostics.

VMEC2000: funct3d.f + residue.f90. cache=None forces a preconditioner/norm/tcon refresh (as VMEC does whenever iter2 == iter1); pass the returned diagnostics.cache back in to reproduce the ns4 = 25 cadence. The returned gc is in the same signed spectral packing as state and feeds vmex.core.step.momentum_update() directly.

vmex.core.solver.solve(source: VmecInput | RunSetup, resolution: Resolution | None = None, *, ftol: float | None = None, max_iterations: int | None = None, mode: str = 'cli', time_step: float | None = None, tcon0: float | None = None, gamma: float | None = None, nstep: int | None = None, lconm1: bool = True, verbose: bool = False, emit=<built-in function print>, initial_state: SpectralState | None = None, device: Any = None, precon_type: str | None = None, prec2d_threshold: float | None = None, prec2d: Prec2DConfig | None = None) → SolveResult¶

Single-grid fixed-boundary solve (VMEC2000 eqsolve.f).

source is a parsed vmex.core.input.VmecInput (recommended; supplies the delt/tcon0/gamma/nstep/ftol/niter defaults, with the keywords overriding) or a prebuilt vmex.core.setup.RunSetup (requires resolution). The resolution defaults to the first ns_array stage (read_indata.f grid rules). Convergence requires fsqr, fsqz, fsql <= ftol simultaneously (evolve.f). mode="cli" runs a Python loop over jitted 10-iteration blocks with host residual checks and VMEC2000-format printing (verbose=True); mode="jit" runs one lax.while_loop over the same traced body.

Returns a SolveResult on convergence. Raises VmecJacobianError when the initial Jacobian changes sign twice (after one guess_axis retry — the eqsolve.f ijacob == 0 path) or at ijacob >= 75 (jac75_flag), and VmecConvergenceError when max_iterations is exhausted (more_iter_flag); both carry the final iteration and (fsqr, fsqz, fsql) diagnostics.

initial_state hot-restarts the solve from a previous SpectralState at the same resolution (e.g. result.state of an earlier solve on a perturbed boundary — VMEC++-style hot restart; use vmex.core.multigrid.interpolate_state() first when ns differs). The R/Z edge row of the provided state is replaced by the input’s processed boundary (the edge never evolves in fixed-boundary mode, so keeping the old row would silently re-solve the old boundary); the interior and lambda are kept.

device places the jitted iteration lanes: "cpu"/"gpu"/ "cuda"/"tpu" or a jax.Device (always honored), or None (default) to apply the measured small-work-to-CPU policy of vmex.core.device — which never overrides a user-pinned JAX_PLATFORMS/JAX_PLATFORM_NAME.

precon_type ("NONE" default) with a finite prec2d_threshold — or an explicit prec2d Prec2DConfig — switches on the optional 2D block preconditioner (VMEC2000 precon2d.f): once fsqr + fsqz + fsql < prec2d_threshold the iteration replaces the 1D radial force direction by a matrix-free Newton step (exact Hessian-vector products via jax.jvp, solved with solvax.gmres()), converging stiff cases (high beta/aspect/mode-number) in far fewer iterations. The default (NONE) path is byte-identical to the 1D-only solver.

vmex.core.solver.hot_restart_state(rt: SolverRuntime, state: SpectralState) → SpectralState¶

Adapt a previous solve’s state to this runtime’s boundary (hot restart).

In fixed-boundary mode the R/Z edge row never evolves, so restarting from state unchanged would silently re-solve the OLD boundary. Replacing only the edge row injects a discontinuous shear between the last two surfaces (measured: initial fsqr ~ 0.5 on cth, i.e. worse than the fresh interior guess). Instead the boundary delta is spread smoothly into the volume with the profil3d.f interior-guess radial profile — sqrts**m for m > 0, linear in s for m = 0 (the axis is held fixed) — which lands the edge row exactly on the new boundary (facj(ns) = 1) and keeps the interior near equilibrium (measured: initial fsqr ~ 4e-6 for a 1% RBC(0,1) perturbation on cth). Lambda is carried over unchanged.

vmex.core.solver.runtime_with_baselines(rt: SolverRuntime, state: SpectralState) → SolverRuntime¶

Rebind rcon0/zcon0 to a new starting state (funct3d.f).

VMEC2000 sets the constraint baselines from the current state whenever iter2 == iter1 — i.e. at the start of every grid. A runtime from prepare_runtime() carries baselines for the profil3d.f interior guess; callers that start from a different state (hot restart via solve(initial_state=...), multigrid stages starting from the interp.f interpolant) must rebind them, or the constraint force — and hence the converged equilibrium — is subtly wrong (observed as a ~1e-8 relative wb shift on the nfp4_QH ladder).

Coarse -> fine radial interpolation of the spectral state (multigrid).

VMEC2000: Sources/TimeStep/interp.f — when the NS_ARRAY ladder moves to the next radial resolution, the converged spectral coefficients xc are interpolated linearly in radius onto the new grid with a VMEC-specific convention (VMEC++ performs the same linear interpolation of the spectral coefficients over the sqrt(s)-scaled internal representation):

scale the coefficients by scalxc (profil3d.f) so odd-m harmonics enter in VMEC’s internal 1/sqrt(s) representation — linear in sqrt(s) near the axis;
extrapolate odd-m modes to the axis on the scaled array, x(js=1) = 2*x(js=2) - x(js=3) (Fortran 1-based);
interpolate linearly between the bracketing coarse surfaces using interp.f’s js1/js2/xint uniform-grid construction;
divide by scalxc on the fine grid to return unscaled (internal physical) coefficients — the state enters and exits WITHOUT the scalxc factor, exactly like the solver’s SpectralState;
zero odd-m coefficients on the output axis row (edge convention sqrts(ns) = 1 is built into scalxc).

The interpolation acts on the m = 1-constrained internal coefficients that vmex.core.solver evolves (interp.f interpolates the internal xc, which is in the constrained basis): every step above is a per-mode linear map that mixes only coefficients with the same poloidal mode number m, so it commutes with the signed-(m, n) packing and with the m = 1 constraint rotation, and no basis conversion is required.

Math ported from the parity-proven legacy port vmex/multigrid.py (interp_vmec_radial_coeffs). Pure JAX, jit-compatible (ns_coarse/ns_fine are static shape information), no host round-trips of traced values.

vmex.core.multigrid.interpolate_coefficients(x_coarse: Any, *, m: ndarray, ns_fine: int) → Any¶

Interpolate one (ns_coarse, mnmax) coefficient array to ns_fine.

VMEC2000 interp.f (see the module docstring for the convention). x_coarse is unscaled (no scalxc); the result is unscaled on the fine grid. m gives the poloidal mode number of each column (static numpy). ns_coarse == ns_fine returns the input unchanged (the legacy short-circuit); the interior-surface values are reproduced exactly in that case by the general path too, since xint vanishes identically.

vmex.core.multigrid.interpolate_state(state_coarse: SpectralState, *, ns_fine: int, modes: ModeTable, ns_coarse: int | None = None) → SpectralState¶

Interpolate a coarse solver state onto a finer radial grid.

VMEC2000 interp.f: the multigrid coarse -> fine transfer of xc between NS_ARRAY stages. state_coarse is the SpectralState of the converged coarse stage — signed-(m, n) internal packing, m = 1-constrained, odd-m WITHOUT the scalxc factor — and the result is in the same representation with ns_fine surfaces, ready for vmex.core.solver.evaluate_forces() on the fine SolverRuntime.

modes must be the mode_table(mpol, ntor) shared by both stages (multigrid only changes ns). ns_coarse is optional (checked against the array shapes when given). Jit-compatible: ns_fine and modes are static, all array work is traced jax.numpy.

vmex.core.multigrid.solve_multigrid(inp: VmecInput, ns_array=None, ftol_array=None, niter_array=None, *, mode: str = 'cli', lconm1: bool = True, verbose: bool = False, emit=<built-in function print>, initial_state: SpectralState | None = None, time_step: float | None = None, tcon0: float | None = None, gamma: float | None = None, nstep: int | None = None, device: Any = None, raise_on_max_iterations: bool = True) → SolveResult¶

Fixed-boundary multigrid solve over the NS_ARRAY ladder.

VMEC2000 Sources/TimeStep/runvmec.f: for each grid igrid the radial resolution nsval = ns_array(igrid) is solved with tolerance ftol_array(igrid) and iteration cap niter_array(igrid); stages with nsval below the best resolution reached so far are skipped (IF (nsval < ns_min) CYCLE — decreasing entries are ignored, equal entries re-run), and each executed stage after the first starts from the interp.f coarse -> fine interpolation (interpolate_state()) of the previous stage’s final state. The time step resets to the input DELT at every stage, and each stage prints its own NS = ... banner (verbose=True, mode="cli").

ns_array/ftol_array/niter_array default to the input’s ladder; when given they are broadcast to a common stage count (shorter ftol/niter arrays repeat their last entry). initial_state seeds the first executed stage (hot restart; must match that stage’s ns).

Intermediate stages are allowed to exhaust their iteration cap (more_iter_flag — VMEC2000 proceeds to the next grid); any other failure raises immediately, and the final stage must converge (VmecConvergenceError otherwise), exactly like vmex.core.solver.solve(). With raise_on_max_iterations=False a final stage that merely hits NITER returns its last state instead (converged=False, ier_flag = more_iter_flag) — VMEC2000’s own behavior, which writes the NITER-exhausted state to the wout file.

Executable reuse: stage runtimes are structural pytrees (solver.py, Phase 2 item (1)), so one XLA executable is compiled per distinct stage structure (ns, ftol, niter, ...) per session, and repeated ladders (parameter scans, hot restarts) recompile nothing. Full radial padding to max(ns_array) — ONE executable for all stages — is the recorded follow-up (§7 item 1); it requires masked radial reductions through geometry/fields/forces/preconditioner and is not attempted here.

device places each stage’s jitted lanes (see vmex.core.solver.solve()): an explicit "cpu"/"gpu"/ jax.Device is always honored; None (default) applies the measured per-stage policy of vmex.core.device (small per-iteration work solves on CPU) unless the user pinned JAX_PLATFORMS.

Returns the final stage’s SolveResult.

Backend (CPU/GPU) selection policy for the core solve lanes.

Measured basis: benchmarks/gpu_baseline.json (2026-07-09, 2x RTX A4000, jax 0.6.2 cuda12) — see its meta.notes and commit a324f503:

Per-iteration throughput favours the GPU at every tested size in the legacy lane (0.83 ms vs 1.90 ms at ns=35, mpol=2, ntor=2, up to 3x on NuhrenbergZille-class decks: 90 s vs 277 s wall).
The GPU pays fixed per-solve overheads (~0.2-0.4 s dispatch/transfer floor plus compile/cache-load on cold processes), so small decks that converge in well under a second of CPU work finish faster on the CPU (solovev: 0.043 s CPU vs 0.289 s CUDA warm wall; cth_like_fixed_bdy: 0.198 s vs 0.383 s).

The rule implemented here uses the per-iteration work proxy ns * mnmax * nznt (radial surfaces x spectral modes x angular grid — the cost driver of the totzsps/tomnsps batched matmuls that dominate one funct3d pass):

deck (first NS_ARRAY stage)	work proxy	warm wall CPU/GPU	winner
solovev (11610)	660	0.043 / 0.289 s	cpu
nfp4_QH_warm_start (35848)	13,440	0.954 / 0.574 s	gpu*
cth_like_fixed_bdy (155324)	24,300	0.198 / 0.383 s	cpu
LandremanPaul2021_QA (16128240)	491,520	14.54 / 4.19 s	gpu
NuhrenbergZille_1988_QHS (11162286)	509,652	276.9 / 90.1 s	gpu

Below GPU_MIN_ITERATION_WORK the measured difference is < 0.5 s either way (the * misclassification costs ~0.4 s); above it the GPU wins by 2-3x and increasingly more with size. The threshold 100_000 sits between the two clusters (geometric mean of 24.3e3 and 491.5e3 ~ 109e3).

The policy is a default only: an explicit device= argument to solve/solve_multigrid always wins, and when the user pinned the JAX platform themselves (JAX_PLATFORMS/JAX_PLATFORM_NAME) the automatic policy stands down entirely (resolve_device() returns None).

vmex.core.device.GPU_MIN_ITERATION_WORK = 100000¶: Minimum ns * mnmax * nznt per-iteration work for the GPU to be the recommended default (see the measured table in the module docstring).

vmex.core.device.iteration_work(resolution: Any) → int¶: Per-iteration work proxy ns * mnmax * nznt of a Resolution.

vmex.core.device.recommended_device(resolution: Any) → str¶

"cpu" or "gpu": the measured-rule recommendation for one stage.

Purely resolution-based (benchmarks/gpu_baseline.json thresholds; see the module docstring); does not check what hardware is present — use resolve_device() for the availability- and pin-aware decision.

vmex.core.device.resolve_device(device: Any, resolution: Any)¶

Map a device= argument to a concrete jax.Device (or None).

None means “leave placement alone” (no jax.default_device wrap):

explicit device ("cpu"/"gpu"/"cuda"/"rocm"/"tpu" or a jax.Device) is always honored — missing hardware raises;
device=None applies recommended_device() unless the user pinned JAX_PLATFORMS/JAX_PLATFORM_NAME (never override an explicit choice), the recommended platform is not available, or the recommendation already matches the default backend.

vmex.core.device.resolve_implicit_device(device: Any, resolution: Any)¶

Device for the implicit-gradient Jacobian / adjoint GMRES (or None).

Unlike the forward solve, the jac="implicit" path builds a per-dof vmapped forward-implicit-differentiation graph — dozens of preconditioned GMRES solves (each with control flow), one per boundary Fourier dof — whose XLA compile grows with the dof count and whose evaluation is kernel-launch bound. Measured on 2x RTX A4000 (R1, benchmarks notes) it is slower on the GPU than on the CPU at every optimization size tested: a max_mode=2 QH stage (24 dofs) did not finish a single Jacobian eval in 37 min on the GPU, versus minutes on the CPU; the forward solve itself is a host callback that never touches the accelerator. So the default here is always the CPU:

an explicit device ("cpu"/"gpu"/… or a jax.Device) is still honored (delegated to resolve_device());
a user JAX_PLATFORMS/JAX_PLATFORM_NAME pin stands down (returns None — leave placement alone);
otherwise pin to the CPU when the default backend is an accelerator, or leave placement untouched (None) when already on the CPU.

resolution is accepted for signature parity with resolve_device() (and in case a size-dependent rule is wanted later); it is unused today.

vmex.core.device.device_context(device: Any, resolution: Any)¶

Context manager placing a solve stage on the resolved device.

Returns jax.default_device(dev) for the resolve_device() result, or a null context when placement should be left untouched.

Free boundary¶

NESTOR vacuum solve (Merkel Green’s-function method), pure JAX.

Computes the scalar magnetic potential on the plasma boundary from the normal component of the external field, following P. Merkel [J. Comp. Phys. 66, 83 (1986)] as implemented by VMEC2000’s NESTOR (Sources/NESTOR_vacuum/{precal,surface,bextern,analyt,greenf,fourp, fouri,scalpot,vacuum}.f).

The math is ported from the legacy parity-proven implementation (vmex.solvers.free_boundary.jax_nestor_operator — host table builders — and ...free_boundary.adjoint.vmec_nestor — the JAX assembly), cleaned into one module:

vacuum_basis() — all geometry-independent tables (VMEC precal.f: mode tables, weighted sin/cos projection bases, cmns analytic-integral coefficients, tan tables, per-period trig tables, fourp index maps). Host NumPy, cached per resolution.
make_vacuum_solver() — jit-compiled closures over a basis:
- full(boundary, bexni): the complete NESTOR update (ivacskip == 0): non-singular Green-function source/kernel (greenf + fourp), analytic singular terms (analyt), mode projection (fouri) and the dense mnpd2 x mnpd2 solve (solver) — returns potvac plus the cached pieces (mode_matrix = amatsav, bvec_nonsing = bvecsav).
- skip(boundary, bexni, bvec_nonsing, mode_matrix): the incremental update (ivacskip != 0): only the analytic source is recomputed and the cached matrix is reused (scalpot.f skip branch).
vacuum_channels() — surface field from potvac: B_u = bexu + d(pot)/du etc., contravariant components through the boundary metric, and bsqvac = |B|^2/2 (vacuum.f tail).

Conventions (identical to VMEC NESTOR):

Angular grid: theta_j = 2*pi*j/nu_full for j < ntheta3 (full range when lasym), zeta_k = 2*pi*k/nzeta per field period.
Rv/Zv and the second derivatives are geometric-phi derivatives (xn = n*nfp in surface.f); the onp = 1/nfp factors below fold them into per-period metric quantities exactly as surface.f does.
bexni = -(B.n) * wint * (2*pi)**2 with the non-unit normal n = signgs*(R*Zu, Ru*Zv - Rv*Zu, -R*Ru) (bextern.f).

class vmex.core.vacuum.VacuumBasis(mf: int, nf: int, nfp: int, nvper: int, lasym: bool, ntheta3: int, nzeta: int, nu_full: int, nuv3: int, nuv_full: int, mnpd: int, mnpd2: int, mn0: int, onp: float, theta: ndarray, zeta: ndarray, wint: ndarray, xmpot: ndarray, n_raw: ndarray, sin_phase: ndarray, cos_phase: ndarray, sinmni: ndarray, cosmni: ndarray, imirr: ndarray, imirr_full: ndarray, cmns: ndarray, idx_all: ndarray, tanu: ndarray, tanv: ndarray, cosuv: ndarray, sinuv: ndarray, cosper: ndarray, sinper: ndarray, cosv_tab: ndarray, sinv_tab: ndarray, cosui: ndarray, sinui: ndarray)¶

Geometry-independent NESTOR tables (host NumPy; precal.f).

Shapes: sin_phase/cos_phase/sinmni/cosmni are (nuv3, mnpd); theta/zeta/wint are flat (nuv3,) over the reduced (ntheta3, nzeta) grid; cmns is (mf+nf+1, mf+1, nf+1).

class vmex.core.vacuum.VacuumBoundary(R: Any, Z: Any, Ru: Any, Zu: Any, Rv: Any, Zv: Any, ruu: Any, ruv: Any, rvv: Any, zuu: Any, zuv: Any, zvv: Any)¶

Boundary geometry on the reduced (ntheta3, nzeta) grid.

Ru/Zu: theta derivatives; Rv/Zv (and all second derivatives): geometric-phi derivatives (xn = n*nfp), exactly as surface.f.

vmex.core.vacuum.build_cmns(*, mf: int, nf: int, onp: float) → ndarray¶: VMEC precal.f cmns(l, m, n) analytic-integral coefficients.

vmex.core.vacuum.make_vacuum_solver(basis: VacuumBasis, *, signgs: int = -1) → VacuumSolver¶: Build the jit-compiled full/skip NESTOR updates for one basis.

vmex.core.vacuum.precal_tan_tables(*, nu: int, nv: int, nvper: int) → tuple[ndarray, ndarray]¶: VMEC precal.f tanu/tanv tables consumed by greenf.f.

vmex.core.vacuum.vacuum_basis(*, mf: int, nf: int, ntheta3: int, nzeta: int, nfp: int, lasym: bool, wint: ndarray) → VacuumBasis¶

Build every geometry-independent NESTOR table for one resolution.

mf = mpol + 1 and nf = ntor (VMEC vacmod0); wint are the VMEC angular integration weights on the (ntheta3, nzeta) grid. Ported from the legacy build_vmec_mode_basis + ensure_vmec_nonsingular_kernel_tables.

vmex.core.vacuum.vacuum_channels(*, basis: VacuumBasis, potvac: Any, bexu: Any, bexv: Any, guu: Any, guv: Any, gvv: Any) → tuple[Any, Any, Any, Any, Any]¶

(bsqvac, bsubu, bsubv, bsupu, bsupv) on the boundary grid.

bexu/bexv are the covariant external-field components with the geometric-phi convention (bexv = Rv*br + R*bp + Zv*bz); guu/guv/ gvv the matching physical surface metric. Equivalent to the vacuum.f tail (whose huv = 0.5*nfp*guv_b and hvv = nfp^2* gvv_b reduce to exactly this physical metric). bsqvac = |B|^2/2.

Free-boundary solve: NESTOR vacuum coupling around the core solver.

Implements the funct3d.f free-boundary block (VMEC2000) on top of the fixed-boundary iteration of vmex.core.solver:

Activation: ivac starts at -1 and increments on every iteration with iter2 > 1 and fsqr + fsqz <= 1e-3; the first vacuum call promotes ivac 0 -> 1 (vacuum.f), prints the In VACUUM block, and triggers the soft-start restart (restart_iter with irst = 2: state <- best stored state, zero velocity, delt *= 0.9, iter1 = iter2, ijacob += 1); eqsolve.f then prints the VACUUM PRESSURE TURNED ON banner and sets ivac = 2.
Cadence: ivacskip = mod(iter2 - iter1, nvacskip) (forced 0 while ivac <= 2); on full steps (ivacskip == 0) the Green-function kernel/matrix is rebuilt and nvacskip = max(nvskip0, 1/max(0.1, 1e11*(fsqr+fsqz))); on skip steps only the analytic source is refreshed against the cached matrix (scalpot.f).
Edge force: bsqvac + presf(ns) enters the R/Z edge force rows via the SolverRuntime free-boundary seam (lfreeb/bsqvac_edge/presf_ns_scale — see solver._evaluate), the edge row is evolved (jmax = ns) and the rcon0/zcon0 constraint baselines are damped by 0.9 per active iteration (funct3d.f).

The iteration itself runs the same traced body as the fixed-boundary lanes. Scheduling (plan item F.2): the pre-activation fixed-boundary iterations run as one jitted lax.while_loop (_preactivation_lane()) and the whole post-turn-on steady state — vacuum cadence, constraint damping, bsqvac_edge refresh and the eqsolve iteration — as another (_make_vacuum_lane()); only the single turn-on pass is host-stepped, because it applies the one-time soft restart and prints the banners. The per-iteration numerics are identical to the per-pass host driver.

Known divergence from VMEC2000 (documented): at turn-on VMEC computes the turn-on iteration’s forces from the pre-restart geometry while evolving the restored state; here the restart is applied before the turn-on iteration, so that iteration’s forces come from the restored state. The golden free-boundary fixture is chaotic/unconverged past turn-on, so trajectories are compared structurally, not pointwise.

class vmex.core.freeboundary.FreeBoundaryState(ivac: int = -1, nvacskip: int = 1, nvskip0: int = 1, turned_on: bool = False, banner_pending: bool = False, delbsq: float = 1.0, bsqvac: ndarray | None = None, mode_matrix: Any = None, bvec_nonsing: Any = None, potvac: ndarray | None = None, ctor: float = 0.0, rbtor: float = 0.0, vacuum_calls: int = 0, full_updates: int = 0)¶: Host cadence state + NESTOR cache (funct3d.f module variables).

vmex.core.freeboundary.boundary_from_coefficients(*, rmnc: ndarray, zmns: ndarray, rmns: ndarray | None, zmnc: ndarray | None, modes: ModeTable, basis: VacuumBasis) → VacuumBoundary¶

Sample the boundary surface on the NESTOR grid (surface.f).

rmnc… are wout-convention edge coefficients over the signed modes table. Angles: theta/zeta from basis (per-period zeta); xn = n*nfp so all v-derivatives are geometric-phi derivatives, exactly as surface.f.

vmex.core.freeboundary.solve_free_boundary(inp: VmecInput, *, mgrid_path: str | Path | None = None, external_field: MgridField | None = None, resolution=None, ftol: float | None = None, max_iterations: int | None = None, verbose: bool = False, emit=<built-in function print>, error_on_no_convergence: bool = True) → SolveResult¶

Single-grid free-boundary solve (eqsolve.f + funct3d.f IVAC0).

external_field overrides the mgrid file (any MgridField-compatible object with a b_cyl(r, phi, z) method — e.g. a direct-coil Biot-Savart field). Raises MgridNotFoundError when the deck’s mgrid file is missing and no field is supplied (callers such as the CLI implement VMEC2000’s warn-and-fall-back-to-fixed-boundary policy).

error_on_no_convergence=False returns the final state instead of raising when NITER is exhausted (useful against unconverged goldens).

Differentiable free-boundary residual via virtual casing (R15.3 + R19).

This module adds a differentiable free-boundary path that complements — and does not touch — the NESTOR forward solve in vmex.core.freeboundary (R15.1/R15.2, VMEC2000-parity). The idea (DESC R17.8): instead of differentiating the NESTOR fixed point, express the free-boundary condition as a smooth objective. At the plasma-vacuum interface the total field is tangent,

B_out . n = 0,

and (finite beta) pressure balance holds,

|B_in|^2 + 2 mu0 p = |B_out|^2,

where B_out = B_coil + B_plasma is the external coil field plus the plasma’s own field, obtained through the virtual-casing principle. Coil / extcur dofs -> B_ext -> boundary residual -> gradients, all in JAX, with no NESTOR-adjoint.

Reuse, not re-implementation¶

The virtual-casing math is reused verbatim from uwplasma/virtual_casing_jax (VirtualCasingExteriorField, VmecSurfaceFieldData, and the accurate on-surface singular-quadrature integral VirtualCasingJAX.compute_internal_B). This module only (a) adapts a converged/trial vmex boundary + total field into the package’s VmecSurfaceFieldData and (b) wires the resulting plasma field to an MgridField or a plain xyz->B callable (e.g. an ESSOS essos.coils.Coils Biot-Savart field) to form the differentiable residual.

Key structural fact that makes this cheap and well-posed: for a fixed trial boundary the plasma’s own field on that boundary does not depend on the coil dofs, so it is precomputed once via the accurate on-surface virtual-casing integral and frozen as a constant. The residual is then a smooth JAX function of the external-field dofs alone — an MgridField’s extcur, or the dofs of a callable coil field (e.g. an ESSOS Coils’ Fourier dofs / currents) — and FD-validates to ~1e-9 (see tests/test_freeboundary_diff.py).

The full single-stage piece — letting the boundary shape dofs vary, so the plasma field itself depends on them through a re-solve — is now supported: surface_field_data_from_state() rebuilds the virtual-casing surface field traceably from a live equilibrium state, so jax.grad threads through the implicit adjoint (boundary) and virtual casing (coils) at once. See examples/single_stage_simultaneous_opt.py and the True single-stage section of docs/optimization.rst.

virtual_casing_jax is an optional dependency (pip install vmex[freeb] or pip install -e /path/to/virtual_casing_jax). Importing this module raises a clear error if it is missing.

vmex.core.freeboundary_diff.MU0 = 1.2566370614359173e-06¶: Vacuum permeability [T m / A] (VMEC2000 mu0 convention).

vmex.core.freeboundary_diff.surface_field_data_from_wout(wout, *, nphi: int = 32, ntheta: int = 32, s_index: int = -1, use_stellsym: bool = True) → VmecSurfaceFieldData¶

Build a VmecSurfaceFieldData from a wout.

The wout may be a converged WoutData (from read_wout() or wout_from_state() — i.e. a trial boundary too). The boundary (last full-mesh surface, s_index=-1) is synthesised on a single field period theta in [0, 2pi), phi in [0, 2pi/nfp):

gamma = (R cos phi, R sin phi, Z),
B_total = B^theta e_theta + B^phi e_phi (Cartesian), with the half-mesh bsup{u,v} edge-extrapolated 1.5 x[-1] - 0.5 x[-2],
normal / area_vector = e_theta x e_phi (oriented outward),

all in the package’s structure-of-arrays layout (3, nphi, ntheta).

The construction is validated by |B_total . n| / |B| ~ 1e-16 on a converged equilibrium (the VMEC free-boundary condition), see the module test.

vmex.core.freeboundary_diff.surface_field_data_from_state(inp, state, *, nphi: int = 32, ntheta: int = 32, s_index: int = -1, use_stellsym: bool = True) → VmecSurfaceFieldData¶

Traceable VmecSurfaceFieldData straight from a SpectralState.

Unlike surface_field_data_from_wout() — which reads a materialised (numpy) wout and so cannot be differentiated in the boundary — this rebuilds the boundary geometry and contravariant-B spectra with the same jnp recipe the wout writer uses (m1_constrained_to_physical(), real_space_geometry(), wout_field_tables()), never leaving the device. The result differentiates in state (hence in the boundary DOFs through the implicit adjoint), which is what makes a simultaneous plasma-boundary + coil single-stage objective possible: jax.grad threads through both this surface field and the coil field.

inp supplies the static resolution / profile metadata; state the (possibly traced) spectral geometry. Stellarator-symmetric only for now (lasym uses the same path but is untested here).

vmex.core.freeboundary_diff.plasma_field_on_boundary(surface_data: VmecSurfaceFieldData, *, digits: int = 4, chunk_size: int = 1024, quad_nt: int | None = None, quad_np: int | None = None, precision=None) → jax.Array¶

Plasma’s own Cartesian field on its boundary via on-surface virtual casing.

Uses the accurate singular-quadrature on-surface integral VirtualCasingJAX.compute_internal_B (the same routine the package’s parity tests exercise), which is well-behaved on the source surface — unlike the off-surface schedule, which is near-singular there. Returns (3, nphi, ntheta) in the same layout as surface_data.B_total.

This is the internal virtual-casing branch (currents inside the LCFS = the plasma current), i.e. the SIMSOPT VirtualCasing.B_external_normal convention: the coils must supply -B_plasma . n for B_out . n = 0.

class vmex.core.freeboundary_diff.FreeBoundaryDiffProblem(gamma: Array, normal: Array, weights: Array, phi_grid: Array, Bn_plasma: Array, B_plasma: Array, Bin_mag2: Array, p_edge: float, nfp: int)¶

A differentiable free-boundary residual for a fixed trial boundary.

Bundles the (constant, coil-dof-independent) boundary geometry and plasma field, precomputed once via virtual casing, with objective methods that are smooth JAX functions of an external field’s dofs. Build with from_wout() (or from_equilibrium()).

gamma, normal: Boundary position and outward unit normal, (3, nphi, ntheta).

weights¶

Area-element weights (surface Jacobian), normalised to sum 1.

Type:: jax.jaxlib._jax.Array

phi_grid¶

Per-toroidal-plane physical angle (nphi,) (for the mgrid path).

Type:: jax.jaxlib._jax.Array

Bn_plasma¶

Plasma’s own normal field B_plasma . n on the boundary (nphi, ntheta) — constant w.r.t. the coil dofs.

Type:: jax.jaxlib._jax.Array

B_plasma¶

Plasma’s own Cartesian field on the boundary (3, nphi, ntheta).

Type:: jax.jaxlib._jax.Array

Bin_mag2¶

|B_total|^2 (internal side) on the boundary (nphi, ntheta).

Type:: jax.jaxlib._jax.Array

p_edge¶

Edge pressure [Pa] for the pressure-balance residual (0 at a true LCFS).

Type:: float

classmethod from_surface_data(surface_data: VmecSurfaceFieldData, *, p_edge: float = 0.0, digits: int = 4, chunk_size: int = 1024, quad_nt: int | None = None, quad_np: int | None = None, precision=None) → FreeBoundaryDiffProblem¶

Precompute the constants (virtual-casing plasma field) from surface data.

Pass precision (from plan_vc_precision(), selected once from a concrete surface) when this runs inside a jax.grad over the boundary geometry, so the virtual-casing plasma field is differentiable in the surface rather than tripping the precision auto-selection’s concretization.

classmethod from_wout(wout, *, nphi: int = 32, ntheta: int = 32, s_index: int = -1, p_edge: float = 0.0, digits: int = 4, chunk_size: int = 1024, quad_nt: int | None = None, quad_np: int | None = None) → FreeBoundaryDiffProblem¶: Build from a converged/trial wout (see surface_field_data_from_wout()).

classmethod from_equilibrium(eq, **kwargs) → FreeBoundaryDiffProblem¶: Build from an Equilibrium (uses eq.wout).

external_B(external_field: Any) → Array¶: External Cartesian field B_ext on the boundary, (3, nphi, ntheta).

external_Bn(external_field: Any) → Array¶: External normal field B_ext . n on the boundary, (nphi, ntheta).

total_B_out(external_field: Any) → Array¶: Vacuum-side total field B_out = B_plasma + B_ext, (3, nphi, ntheta).

bnormal_residual(external_field: Any) → Array¶

Free-boundary normal residual (B_plasma + B_ext) . n, (nphi, ntheta).

Zero for a self-consistent free-boundary equilibrium; equals B_ext . n - B_external_normal in the SIMSOPT stage-2 convention.

pressure_balance_residual(external_field: Any) → Array¶: Pressure-balance residual |B_out|^2 - |B_in|^2 - 2 mu0 p, (nphi, ntheta).

bnormal_objective(external_field: Any) → Array¶: Area-weighted mean square of bnormal_residual() (a scalar).

pressure_balance_objective(external_field: Any) → Array¶: Area-weighted mean square of pressure_balance_residual() (a scalar).

objective(external_field: Any, *, pressure_weight: float = 0.0) → Array¶: Combined residual J_bn + pressure_weight * J_pressure (a scalar).

vmex.core.freeboundary_diff.external_B_cartesian(external_field: Any, gamma: Array, phi_grid: Array | None = None) → Array¶

Cartesian external field at boundary points gamma = (3, nphi, ntheta).

Dispatches on the external-field type, staying differentiable in its dofs:

MgridField -> trilinear mgrid (diff. in extcur),
a plain callable xyz(..., 3) -> B(..., 3) (e.g. an ESSOS Coils Biot-Savart field, lambda pts: coils.B(pts); diff. in its own dofs).

Returns (3, nphi, ntheta).

vmex.core.freeboundary_diff.have_virtual_casing_jax() → bool¶: Return True iff virtual_casing_jax is importable.

MAKEGRID mgrid file IO and differentiable field interpolation.

This module is the clean-core home of the VMEC free-boundary external-field inputs (§8):

MgridData — an immutable snapshot of a VMEC2000/MAKEGRID mgrid netCDF file (grid extents, dimensions, per-coil-group cylindrical field tables, coil-group labels, raw currents).
read_mgrid() / write_mgrid() — netCDF round-trip in the exact MAKEGRID layout consumed by VMEC2000’s read_mgrid (NETCDF3 classic, br_001/bp_001/bz_001… variables on (phi, zee, rad)).
MgridField — a JAX pytree wrapping the field tables plus external currents, providing extcur-scaled trilinear interpolation of B(r, phi, z) that is jit- and grad-compatible.

VMEC2000 counterpart: Sources/NESTOR_vacuum/mgrid_mod.f (read_mgrid, becoil). The IO layer is ported from the legacy parity-proven vmex.solvers.free_boundary.mgrid; the interpolation kernel from vmex.external_fields.mgrid_jax.

ESSOS compatibility (`essos.mgrid.MGrid`, PR#33)¶

The netCDF layout is identical (same variable names, dimensions, and (phi, zee, rad) field ordering), so files written by either package are readable by both. Known convention divergences:

Label padding: ESSOS centers coil-group names in 30 characters and replaces spaces with underscores (SIMSOPT to_mgrid convention); MAKEGRID left-justifies and pads with blanks. write_mgrid() follows MAKEGRID (blank padding); read_mgrid() strips trailing whitespace so either style reads back cleanly (underscore-padded ESSOS names keep their underscores, exactly as ESSOS’ own reader returns them).
stringsize: ESSOS always writes stringsize = 30; write_mgrid() preserves the label width of the data being written (min 30) so round-trips of MAKEGRID files with longer labels do not truncate.
mgrid_mode: ESSOS’ writer hard-codes "N" and its reader exposes the per-group sums br/bp/bz and unit raw_coil_cur; MgridData keeps the mode and raw currents from the file and leaves scaling to MgridField via extcur.

class vmex.core.mgrid.MgridData(rmin: float, rmax: float, zmin: float, zmax: float, ir: int, jz: int, kp: int, nfp: int, nextcur: int, mgrid_mode: str, coil_groups: tuple[str, ...], raw_coil_cur: tuple[float, ...], br: ndarray, bp: ndarray, bz: ndarray)¶

Contents of a VMEC2000/MAKEGRID mgrid netCDF file.

Field arrays are stored per coil group with shape (nextcur, kp, jz, ir) — the netCDF per-group (phi, zee, rad) layout stacked along a leading coil-group axis. Fields are unscaled coil-group responses; the physical field is sum_g extcur[g] * B_g (mgrid_mode="S", scaled per unit current) or the raw per-group field (mgrid_mode="R"/”N”, raw currents baked in — VMEC then divides extcur by raw_coil_cur).

VMEC2000 counterpart: the module variables filled by read_mgrid in Sources/NESTOR_vacuum/mgrid_mod.f.

rmin: float¶: Grid extents (meters), inclusive at both ends for R and Z.

ir: int¶

ir radial, jz vertical, kp toroidal planes per period.

Type:: Grid dimensions

nfp: int¶: Number of field periods; the phi grid spans [0, 2*pi/nfp), endpoint excluded (plane kp would repeat plane 0 of the next period).

nextcur: int¶: Number of external coil groups (independently scalable currents).

mgrid_mode: str¶: “S” (scaled, per unit current), “R” (raw currents), or “N” (none/raw).

coil_groups: tuple[str, ...]¶: Coil-group labels, trailing whitespace stripped (len == nextcur).

raw_coil_cur: tuple[float, ...]¶: Currents the file was computed with (len == nextcur).

br: ndarray¶: Cylindrical field tables, shape (nextcur, kp, jz, ir) each.

class vmex.core.mgrid.MgridField(br: Any, bp: Any, bz: Any, extcur: Any, rmin: float, rmax: float, zmin: float, zmax: float, nfp: int)¶

Extcur-scaled trilinear mgrid field B(r, phi, z) (pure JAX).

A pytree: the field tables and extcur are differentiable leaves; grid extents and nfp are static metadata, so instances can flow through jax.jit/jax.grad (differentiable w.r.t. field values and extcur everywhere, and w.r.t. coordinates away from cell boundaries).

VMEC2000 counterpart: becoil in Sources/NESTOR_vacuum/mgrid_mod.f (which samples the kp planes directly; this class interpolates periodically in phi, matching the legacy generic path).

classmethod from_mgrid_data(data: MgridData, extcur: Any | None = None) → MgridField¶: Build a field from MgridData; defaults extcur to raw currents.

classmethod from_file(path: str | Path, extcur: Any | None = None) → MgridField¶: Read path (raising MgridNotFoundError if missing) and build a field.

b_cyl(r: Any, phi: Any, z: Any) → tuple[Any, Any, Any]¶: Return (B_r, B_phi, B_z) at cylindrical points (broadcastable).

vmex.core.mgrid.read_mgrid(path: str | Path) → MgridData¶

Read a VMEC2000/MAKEGRID mgrid netCDF file.

Raises:: MgridNotFoundError – If the file does not exist or is not a readable netCDF dataset (host-side check; nothing here is traced).

vmex.core.mgrid.write_mgrid(path: str | Path, data: MgridData) → None¶

Write data as a MAKEGRID-layout netCDF file (NETCDF3 classic).

Dimension and variable names match MAKEGRID’s write_mgrid_nc (and ESSOS MGrid.write): scalars ir/jz/kp/nfp/nextcur/rmin/zmin/rmax/ zmax, char mgrid_mode(dim_00001) and coil_group( external_coil_groups, stringsize), raw_coil_cur(external_coils), and per-group br_XXX/bp_XXX/bz_XXX(phi, zee, rad).

Divergence from ESSOS: coil-group labels are blank-padded (MAKEGRID style) rather than underscore-padded/centered, mgrid_mode and raw_coil_cur are taken from data instead of hard-coded "N"/ones, and coil_group is always written 2-D (ESSOS writes a 1-D char vector when there is a single group; both forms are accepted by read_mgrid(), ESSOS, and VMEC2000).

Differentiation and optimization¶

Implicit differentiation of the fixed-boundary equilibrium (§6).

The converged equilibrium is a root of F(x, p) = 0 with x the SpectralState and p the differentiable run parameters (ImplicitParams: dense INDATA boundary arrays, phiedge, pres_scale, curtor and the am/ai/ac profile coefficients). solve_implicit() wraps the opaque host solver (vmex.core.solver.solve() / solve_multigrid) in jax.custom_vjp; the backward pass solves the adjoint linear system

(dF/dx)^T lambda = g_x

matrix-free (one jax.vjp linearization of the residual, re-applied by a recycling GCROT(m, k) Krylov solve — see _adjoint_solve_gcrot) and returns g_p - lambda^T dF/dp with one more VJP — O(1) memory in the forward iteration count.

Residual formulation (documented choice)¶

F is the self-consistently preconditioned force gc of a single fresh evaluate_forces() pass (cache=None): the bcovar.f preconditioner/force norms/tcon are recomputed from the current (x, p) rather than frozen at the converged cache. This makes F a fixed, smooth function of (x, p) — required by the implicit function theorem — while remaining exactly as well-conditioned as VMEC’s own preconditioned iteration. Correctness: gc = M(x, p) f(x, p) with f the raw (scalxc-scaled) spectral force and M the invertible linear 1D-preconditioner map (scale_m1 + scalfor tridiagonal solves + faclam). At the root, dF = M df + dM f = M df up to O(|f|) = O(ftol), so the implicit gradients of the preconditioned residual equal those of the raw force residual to solver accuracy, and GMRES on dF/dx inherits the preconditioning for free. The raw-force formulation (formulation="raw") is kept for the informational with/without- preconditioner comparison in the tests.

The m=1 constrained Z force is evaluated in its converged branch (zeroed — residue.f90 zeroes gcz(m=1) once fsqz < 1e-6, which always holds at the fixed point), so the corresponding constrained combinations are not degrees of freedom: they are frozen at their converged values, exactly mirroring the forward solver’s behavior near convergence.

Degrees of freedom / boundary handling¶

In fixed-boundary mode the R/Z edge spectral row never evolves: the full state is assembled as x = mask*z + edge_mask*boundary(p) + frozen where z are the evolved dofs, the edge row comes (differentiably) from the boundary parameters, and the remaining entries (structurally zero families, released m=1 combinations, the lambda axis row overwritten by the totzsp closure) are frozen constants. The dof mask is computed once per forward solve from the exact structural zero patterns of gc (row support) and of the x-dependence of gc (column support, one VJP with a random cotangent) at a generically perturbed state — see _dof_mask.

Gradient checking solver-sensitive metrics¶

The adjoint gradient is the derivative of the fixed point of the frozen residual F — the preconditioner/tcon/m=1 branch/dof mask are captured once at the base parameters, not re-derived. For a smooth bulk integral (wb, aspect) a naive central FD through the full host solver already matches jax.grad to rtol <= 1e-6. But a solver-sensitive metric — iota (derived from the current-constrained chips at ncurr=1), the mirror ratio, the magnetic well, the Boozer/QI residual — reads the converged state directly, and a naive re-solve at p ± h lets that convergence logic re-form slightly differently on each side, an O(1) path perturbation that can sign-flip the FD (d(iota_edge)/d(RBC(-1,1)) on li383_low_res: adjoint -0.773, naive FD +0.045). The naive FD is therefore not a valid reference for these metrics; frozen_path_directional_fd() provides the correct one (Newton-solve the frozen F at p ± h), and it reproduces the adjoint to solver accuracy — see tests/test_implicit_grad.py.

Zero-crash typed errors through the callback¶

The forward solve runs behind jax.pure_callback, which converts any host exception into an opaque jax.errors.JaxRuntimeError whose message embeds the whole host traceback (measured ~3.7 KB) and loses the typed exception (__cause__ is None) — breaking the vmex.core.errors zero-crash taxonomy. The relay: _host_solve_and_mask catches any VmecError, stashes it in the module-level _HOST_ERROR slot (host callbacks are serialized per solve, so a single slot suffices) and re-raises a SHORT sentinel RuntimeError; the pure_callback call sites (_callback_solve) catch the runtime error, pop the slot and re-raise the ORIGINAL typed exception with from None to suppress the callback noise. im.run / solve_implicit() therefore fail with a short VmecConvergenceError / VmecJacobianError. Under jax.jit the error instead surfaces at the jit boundary (where the optimize.least_squares zero-crash penalty lanes catch it); the sentinel keeps even that message short.

Parameter map¶

runtime_from_params rebuilds every p-dependent RunSetup field traceably: the readin.f boundary processing (with the theta-flip decision frozen from the reference input — it is discrete), the profil1d.f flux/pressure/current profiles, and the funct3d.f constraint baselines rcon0/zcon0 (which depend on the boundary only — the edge row of any admissible state). Its output is verified against run_setup() in tests/test_implicit_grad.py.

class vmex.core.implicit.ImplicitParams(rbc: Any, rbs: Any, zbc: Any, zbs: Any, phiedge: Any, pres_scale: Any, curtor: Any, am: Any, ai: Any, ac: Any)¶

Differentiable run parameters (a JAX pytree).

rbc/rbs/zbc/zbs are the dense INDATA boundary arrays, shape (2*ntor + 1, mpol) indexed [n + ntor, m] (physical, un-processed — exactly VmecInput layout, so e.g. RBC(0, 1) is rbc[ntor, 1]). am/ai/ac are the dense profile coefficient arrays; phiedge/pres_scale/curtor scalars.

class vmex.core.implicit.ImplicitConfig(inp: VmecInput, resolution: Resolution, ftol: float, max_iterations: int, mode: str = 'cli', multigrid: bool = False, lconm1: bool = True, adjoint_tol: float = 1e-11, adjoint_restart: int = 30, adjoint_maxiter: int = 300, adjoint_gcrot_m: int = 100, adjoint_gcrot_k: int = 20, hot_restart: bool = False)¶

Static (non-differentiable) context of one implicit solve.

adjoint_gcrot_m: int = 100¶: reverse-adjoint GCROT(m, k) recycling solve (see _adjoint_solve_gcrot): the (dF/dz)^T lambda = g_x solve of the backward pass uses GCROT with this inner FGMRES cycle size m and k deflation directions instead of plain restarted GMRES. Measured (padded nfp4_QH, ns=25): at high mode number restarted GMRES stalls short of adjoint_tol within its budget (the truncated Krylov cycle loses the small eigendirections), while GCROT(100, 20) converges to the same lambda in ~40% fewer matvecs and ~30% less wall — a strictly faster path to the identical converged adjoint.

hot_restart: bool = False¶: seed repeated host solves from the last converged state of this config (optimization trials; the fixed point — hence the gradient — is unchanged, only the iteration count drops). Makes the callback stateful across calls, so keep False for one-shot/diagnostic use.

class vmex.core.implicit.ImplicitSolution(state: SpectralState, wb: Any, wp: Any, wmhd: Any, volume: Any, aspect: Any, iota_axis: Any, iota_edge: Any, runtime: SolverRuntime | None = None)¶

Differentiable outputs of run() (a JAX pytree).

runtime is the SolverRuntime that run() built internally (runtime_from_params(params, cfg)), so objective callers can evaluate further (state, runtime) scalar targets without rebuilding it per evaluation. It is deliberately not part of the pytree (registered as a dropped field): the solution’s established pytree structure — six state leaves plus seven scalars — is unchanged, and a solution that round-trips through flatten/unflatten (e.g. across a jax.jit boundary) comes back with runtime = None. Inside a jax.grad/jax.value_and_grad trace of run() the attribute is available and fully traced, so gradients flow through runtime-consuming objectives exactly as through an explicit runtime_from_params rebuild.

vmex.core.implicit.params_from_input(inp: VmecInput) → ImplicitParams¶

Extract the differentiable parameters of an input as a pytree.

On an accelerator box the pytree is committed to the CPU (vmex.core.device.resolve_implicit_device()): every eager op of a jax.grad/jax.jacrev over run() then executes there, which is where the launch-bound implicit adjoint is fastest — measured 57 s (GPU) vs seconds (CPU) for one solovev value_and_grad (R24). A user JAX_PLATFORMS pin or an already-CPU backend stands the pin down; optimize.least_squares applies the same rule to its dof vector.

vmex.core.implicit.input_with_params(inp: VmecInput, params: ImplicitParams) → VmecInput¶: Host-side: a new VmecInput with the parameter values applied.

vmex.core.implicit.runtime_from_params(params: ImplicitParams, cfg: ImplicitConfig) → SolverRuntime¶

Differentiable (traceable) map p -> SolverRuntime.

Rebuilds every p-dependent RunSetup field with jnp operations: the processed boundary, the profil1d.f flux/mass/current profiles (through vmex.core.setup.flux_profiles(), which is traced in phiedge/pres_scale/curtor/am/ai/ac and in r00), the profil3d.f interior guess (whose edge row is the boundary — the initial interior is an initializer only) and the constraint baselines rcon0/zcon0 (functions of the edge row alone). All p-independent fields (grids, scalxc, axis arrays, static metadata) come from the reference runtime.

vmex.core.implicit.make_config(inp: VmecInput, *, ns: int | None = None, ftol: float | None = None, max_iterations: int | None = None, mode: str = 'cli', multigrid: bool = False, lconm1: bool = True, adjoint_tol: float = 1e-11, adjoint_restart: int = 30, adjoint_maxiter: int = 300, adjoint_gcrot_m: int = 100, adjoint_gcrot_k: int = 20, hot_restart: bool = False) → ImplicitConfig¶: Build the static config; resolution is the (final-stage) grid.

vmex.core.implicit.solve_implicit_with_aux(params: ImplicitParams, cfg: ImplicitConfig)¶: Return (state, dof_mask) using the same callback as solve_implicit.

vmex.core.implicit.implicit_state_pullback_multi_rhs(params: ImplicitParams, cfg: ImplicitConfig, x_star: SpectralState, dof_mask: SpectralState, gbar_batch: SpectralState) → ImplicitParams¶

Batched state-cotangent pullback with shared implicit-linearization setup.

This preserves the scalar solve_implicit VJP and only adds a helper for callers that already have several state cotangents for the same fixed point. It reuses the residual/projector/VJP setup once, then applies the existing single-RHS GMRES per row.

vmex.core.implicit.run(source: VmecInput | str, params: ImplicitParams | None = None, *, ns: int | None = None, ftol: float | None = None, max_iterations: int | None = None, mode: str = 'cli', multigrid: bool = False, lconm1: bool = True, adjoint_tol: float = 1e-11, adjoint_restart: int = 30, adjoint_maxiter: int = 300, adjoint_gcrot_m: int = 100, adjoint_gcrot_k: int = 20) → ImplicitSolution¶

Differentiable fixed-boundary equilibrium: input -> outputs pytree.

params defaults to params_from_input(); pass a perturbed / traced ImplicitParams to differentiate:

inp = VmecInput.from_file("input.solovev")
p0 = params_from_input(inp)
grad = jax.grad(lambda p: run(inp, p).wb)(p0)

wmhd follows the printed WMHD normalization; gamma = 1 inputs get wmhd = nan (as in VMEC). All outputs are differentiable in params (state via the implicit adjoint; scalars additionally through their explicit parameter dependence).

The returned solution also carries the internally built SolverRuntime as sol.runtime (a non-pytree convenience attribute, see ImplicitSolution), so objective code can evaluate additional (state, runtime) targets — e.g. optimize.mean_iota(sol.state, sol.runtime) — without repeating runtime_from_params(params, make_config(...)) per evaluation.

vmex.core.implicit.mhd_energy(state: SpectralState, rt: SolverRuntime) → tuple[Any, Any]¶: (wb, wp) in the wout normalization (bcovar.f), differentiable.

vmex.core.implicit.plasma_volume(state: SpectralState, rt: SolverRuntime) → Any¶

Plasma volume volume_p [m^3] (= (2 pi)^2 * hs * sum vp).

Quadrature note (Item I.7): this is the implicit module’s historical differential-volume sum, pinned by ImplicitSolution.volume and the FD-cached gradient tables of tests/test_implicit_grad.py. The canonical wout-parity boundary quadrature of the same scalar is vmex.core.statephysics.volume() (re-exported as optimize.volume); the two agree to quadrature resolution.

vmex.core.implicit.aspect_ratio(state: SpectralState, rt: SolverRuntime, *, ntheta: int = 128, nzeta: int = 32) → Any¶

VMEC-convention aspect ratio Rmajor_p / Aminor_p (differentiable).

Aminor_p = sqrt(<cross-section area>_zeta / pi) with the area from the shoelace integral -oint Z dR/dtheta dtheta on the boundary, and Rmajor_p = volume_p / (2 pi^2 Aminor_p^2) (aspectratio.f).

Quadrature note (Item I.7): this is the implicit module’s historical shoelace-on-a-fresh-grid variant, pinned by ImplicitSolution.aspect and the FD-cached solovev gradient table of tests/test_implicit_grad.py. The canonical wout-parity aspectratio.f boundary quadrature (internal-grid wint weights, equal to the wout aspect scalar) is vmex.core.statephysics.aspect_ratio() (re-exported as optimize.aspect_ratio); the two agree to quadrature resolution.

vmex.core.implicit.iota_profile(state: SpectralState, rt: SolverRuntime) → Any¶

Full-mesh iotaf (add_fluxes.f90), differentiable.

ncurr = 0: the prescribed profile (p-dependent through ai); ncurr = 1: reconstructed from the converged current-constrained chips exactly as in the solver’s result assembly.

vmex.core.implicit.iota_edge(state: SpectralState, rt: SolverRuntime) → Any¶

Boundary rotational transform iotaf[-1] (differentiable).

Naming note (Item I.7): the same physical scalar as vmex.core.statephysics.edge_iota() (optimize.edge_iota) — identical for ncurr = 1; at ncurr = 0 this evaluates the prescribed full-mesh iotaf endpoint while the wout-parity version extrapolates the half-mesh iotas. edge_iota is provided as an alias here so either spelling works in either module.

vmex.core.implicit.edge_iota(state: SpectralState, rt: SolverRuntime) → Any¶

Boundary rotational transform iotaf[-1] (differentiable).

vmex.core.implicit.residual_fn(cfg: ImplicitConfig, frozen: SpectralState, dof_mask: SpectralState, formulation: str = 'preconditioned') → Callable¶

Return the implicit residual F(z, params) -> masked force pytree.

formulation="preconditioned" (default, used by the adjoint): the self-consistently preconditioned gc of a fresh evaluate_forces() pass — see the module docstring for why its implicit gradients coincide with the raw formulation. formulation="raw": the un-preconditioned (scalxc-scaled, m=1 rotated/zeroed) spectral force — same root, same gradients, but the adjoint GMRES then runs without preconditioning (diagnostic only).

vmex.core.implicit.adjoint_matvec(cfg: ImplicitConfig, params: ImplicitParams, x_star: SpectralState, dof_mask: SpectralState, formulation: str = 'preconditioned') → Callable¶: v -> (dF/dz)^T v for tests/diagnostics (both formulations).

vmex.core.implicit.frozen_path_directional_fd(params: ImplicitParams, cfg: ImplicitConfig, metric_fn: Callable[[SpectralState, SolverRuntime], Any], tangent: ImplicitParams, *, h: float = 0.0001, newton_steps: int = 20, newton_rtol: float = 1e-11) → tuple[float, dict]¶

Central FD of metric_fn along tangent on the frozen solve path.

The correct finite-difference reference for solver-sensitive metrics – iota (derived from the current-constrained chips at ncurr=1), the mirror ratio, the magnetic well, the Boozer/QI residual – whose value reads the converged solver state directly rather than through a smooth bulk integral (wb, aspect, for which a naive re-solve FD is already exact and jax.grad() matches it to rtol <= 1e-6).

A naive full re-solve at params +/- h*tangent lets the solver’s internal convergence logic – the bcovar preconditioner, the tcon constraint scaling, the m=1 gcz zeroing branch (residue.f90), the dof mask, the multigrid schedule, and exactly where the ftol crossing lands – re-form slightly differently at each perturbed point. For a solver-sensitive metric that path variation is an O(1) contribution that can inflate or even sign-flip the finite difference (measured on li383_low_res: d(iota_edge)/d(RBC(-1,1)) = -0.773 from the adjoint, but the naive central FD reads +0.045 – wrong sign).

The implicit adjoint deliberately does not differentiate through that logic: it linearizes the fixed point of the frozen residual F (the preconditioner / mask / branch captured once at params; see the module docstring), which is the stable, physical gradient. This helper reproduces exactly that path – it captures F once at params and Newton-solves F(z, params +/- h*tangent) = 0 (matrix-free, the same linearization the adjoint uses) from the converged z* before central-differencing metric_fn. The result therefore equals jax.grad() of the metric contracted with tangent to solver accuracy – the gradient check a naive re-solve FD cannot provide for these metrics.

Returns (fd, info) where info['newton_res'] are the two frozen-solve residual norms; confirm they are small (an unconverged frozen solve invalidates the comparison).

Optimization objectives and least-squares driver for the new core (§5.1, §10).

Simsopt-style vocabulary for the QA/QH/QP/QI examples on the pure new core:

QuasisymmetryRatioResidual — the two-term quasisymmetry ratio residual of Landreman & Paul (simsopt QuasisymmetryRatioResidual), evaluated from the wout-engine field tables of a converged core state; math ported verbatim from the parity-proven legacy vmex/quasisymmetry.py (quasisymmetry_ratio_residual_from_wout).
practical scalar targets — aspect_ratio(), mean_iota(), edge_iota(), mirror_ratio(), volume(), magnetic_well() — each a pure function of (SpectralState, SolverRuntime) on vmex.core.geometry / vmex.core.fields.
a distilled Goodman-style QI residual (quasi_isodynamic_residual()) keeping exactly the four terms the legacy minimal-seed QI examples exercised (now examples/optimization/QI_optimization.py): level-set bounce-width variance, branch trapped-well width variance, field-line profile consistency, and the branch-shuffle profile comparison. The unused legacy knobs (aligned_profile_*, weighted_shuffle_*, shuffle_profile_nphi_out) were dropped.
least_squares() — a thin scipy.optimize.least_squares() driver over boundary Fourier dofs (pack_boundary()/unpack_boundary()), taking simsopt-style (callable, target, weight) terms.

Helicity conventions (match legacy/simsopt exactly)¶

The QS residual keeps the |B| spectrum aligned with the single helicity chi = helicity_m * theta - helicity_n * nfp * phi — helicity_n is in units of nfp (the internal target mode number is nn = helicity_n * nfp):

QA: (helicity_m, helicity_n) = (1, 0)
QH: (1, -1) (i.e. chi = theta + nfp*phi; legacy/simsopt sign — the plan’s “n = -nfp” written in physical toroidal mode numbers)
QP: (0, 1)

Gradient modes¶

least_squares() defaults to scipy finite differences (jac=None -> "2-point"). jac="implicit" uses the Phase-6 implicit-gradient path (vmex.core.implicit): each trial boundary is solved once through solve_implicit() (a jax.custom_vjp around the host solver) and the exact residual Jacobian comes from forward implicit differentiation of the fixed point — one preconditioned GMRES per boundary dof (a few dozen residual linearizations each) instead of one full equilibrium solve per dof. In implicit mode every objective term must be a traceable function of (SpectralState, SolverRuntime); vector-valued terms exposing a residuals_state method (QuasisymmetryRatioResidual) contribute their full pointwise residual vector, matching the finite-difference stacked-residual cost and Gauss-Newton geometry (internal-grid sampling instead of the wout grid). Wout-engine terms (d_merc(), l_grad_b(), the Boozer-based QI residual) run on host NumPy and are finite-difference-only — l_grad_b_state() is the traceable (state, rt) lane of the same L_grad_B convention for jac="implicit". The implicit parameter map supports lasym via the four RBC/ZBS/RBS/ZBC boundary families and a traceable readin.f delta rotation (FD-validated), so jac="implicit" handles lasym = True as well (the QS-ratio traceable term is symmetric-only).

Measured cost (2026-07-10, RTX A4000, nfp2 circular seed, QS + aspect + iota objective, max_mode=2 -> 24 dofs): warm implicit Jacobian 2.5 s (~1.5 hot-restart equilibrium-solve equivalents, independent of the dof count) vs the 2-point FD Jacobian’s 24 hot solves ~ 39 s — 15.7x; the gap widens linearly with max_mode.

class vmex.core.optimize.Equilibrium(inp: VmecInput, state: SpectralState, runtime: SolverRuntime, result: SolveResult)¶

A converged fixed-boundary equilibrium plus its evaluation contexts.

Objective callables in least_squares() receive one of these. The solver-native pieces (state, runtime) feed the differentiable scalar targets; wout (built lazily, host NumPy) feeds the wout-table objectives (QS ratio residual, Boozer-based QI residual).

property wout: WoutData¶: Full wout dataset of this state (vmex.core.wout, cached).

vmex.core.optimize.solve_equilibrium(inp: VmecInput, *, initial_state: SpectralState | None = None, raise_on_max_iterations: bool = False, **solve_kwargs) → Equilibrium¶

Converge inp with the core multigrid solver -> Equilibrium.

raise_on_max_iterations=False by default: during optimization a NITER-exhausted trial state is still a usable (penalized) sample — VMEC2000 behaves the same way. Extra keywords go to vmex.core.multigrid.solve_multigrid().

class vmex.core.optimize.QuasisymmetryRatioResidual(surfaces, helicity_m: int = 1, helicity_n: int = 0, *, weights: Iterable[float] | None = None, ntheta: int = 63, nphi: int = 64)¶

Two-term quasisymmetry ratio residual (simsopt convention).

On each requested surface the field is sampled on a uniform (theta, phi) grid (VMEC angles) and the pointwise residual

f = [(B x grad B . grad psi)(nn - iota*m) - (B . grad B)(m*G + nn*I)] / B^3

(m = helicity_m, nn = helicity_n * nfp, G/I the Boozer covariant field averages bvco/buco) is weighted by the flux-surface measure sqrt(nfp*dtheta*dphi*|sqrt g| / V') so that total = sum(residuals**2) is simsopt’s surface-averaged QS ratio. f vanishes identically iff |B| depends on the angles only through helicity_m*theta - nn*phi.

The evaluation consumes the parity-proven wout-engine tables (bmnc/gmnc/bsub*/bsup*, vmex.core.nyquist) of a WoutData — from wout_from_state() or any wout_*.nc — ported from legacy quasisymmetry_ratio_residual_from_wout (A/B bit-exact).

compute(wout) → dict[str, Any]¶: Full diagnostics dict from a wout-like object or Equilibrium.

residuals(wout) → Array¶: Flat least-squares residual vector (target 0, weight applied by the driver).

profile(wout) → Array¶: Per-surface sum of squared residuals.

total(wout) → Any¶: Scalar QS ratio objective sum(residuals**2).

J(eq: Equilibrium) → Array¶: Objective-term entry point for least_squares() (residual vector).

residuals_state(state: SpectralState, rt: SolverRuntime) → Array¶

Traceable flat residual vector with sum(r**2) = total_state.

The internal-grid analogue of residuals() (wout tables): the pointwise weighted residual of _pointwise_state() scaled by the square roots of the surface coefficients — this is the residual vector jac="implicit" optimizes, giving the least-squares driver the full pointwise Gauss-Newton geometry.

profile_state(state: SpectralState, rt: SolverRuntime) → Any¶

Traceable weighted per-surface QS totals at surfaces.

weights * interp(surfaces, <f^2> profile) from _pointwise_state(); sum = total_state.

total_state(state: SpectralState, rt: SolverRuntime) → Any¶: Traceable scalar QS objective: sum(profile_state) (see there).

vmex.core.optimize.aspect_ratio(state: SpectralState, rt: SolverRuntime) → Any¶

VMEC aspect ratio Rmajor_p / Aminor_p (aspectratio.f convention).

Aminor_p = sqrt(<cross-section area> / pi), Rmajor_p = volume_p / (2 pi <area>) from the boundary surface quadrature; equals the wout aspect scalar of the same state. This is the canonical (wout-parity) implementation, re-exported as vmex.core.optimize.aspect_ratio; vmex.core.implicit.aspect_ratio() is the implicit module’s historical shoelace-quadrature variant of the same scalar (see there).

vmex.core.optimize.mean_iota(state: SpectralState, rt: SolverRuntime) → Any¶

Mean rotational transform over the half-mesh surfaces (axis excluded).

Matches the legacy optimization mean_iota convention (mean(iotas[1:]), i.e. the mean of the wout iotas profile).

vmex.core.optimize.edge_iota(state: SpectralState, rt: SolverRuntime) → Any¶

Rotational transform at the boundary (wout iotaf[-1] convention: linear extrapolation of the half mesh, 1.5 iotas[-1] - 0.5 iotas[-2]).

Naming note (Item I.7): optimize.edge_iota and vmex.core.implicit.iota_edge() are the same physical scalar — identical for ncurr = 1 (both reconstruct iota from the converged chips); at ncurr = 0 this wout-parity version extrapolates the prescribed half-mesh iotas while the implicit variant evaluates the prescribed full-mesh iotaf endpoint directly. iota_edge is provided as an alias here (and edge_iota in implicit) so either spelling works in either module.

vmex.core.optimize.iota_edge(state: SpectralState, rt: SolverRuntime) → Any¶

Rotational transform at the boundary (wout iotaf[-1] convention: linear extrapolation of the half mesh, 1.5 iotas[-1] - 0.5 iotas[-2]).

vmex.core.optimize.mirror_ratio(state: SpectralState, rt: SolverRuntime, *, s_index: int = -1) → Any¶

Mirror ratio (Bmax - Bmin) / (Bmax + Bmin) on one half-mesh surface.

|B| is evaluated on the solver’s internal angular grid from the half-mesh field state (|B|^2 = 2 (bsq - p), bcovar.f); s_index selects the half-mesh surface (default: outermost). Hard max/min — smooth almost everywhere, adequate for finite-difference least squares (the legacy VMECMirrorRatio softmax knobs were an optimizer nicety only).

vmex.core.optimize.volume(state: SpectralState, rt: SolverRuntime) → Any¶

Plasma volume volume_p [m^3] (wout convention, boundary quadrature).

Canonical (wout-parity) implementation, re-exported as vmex.core.optimize.volume; vmex.core.implicit.plasma_volume() is the implicit module’s sum(vp) variant of the same scalar (see there).

vmex.core.optimize.magnetic_well(state: SpectralState, rt: SolverRuntime) → Any¶

VMEC/simsopt magnetic-well proxy (V'(0) - V'(1)) / V'(0).

V' = dV/ds endpoints are linear extrapolations of the half-mesh differential volume vp (bcovar.f); positive values mean a favorable well (vacuum_well in simsopt). Ported from legacy vmex.finite_beta.magnetic_well_from_vp.

vmex.core.optimize.d_merc(eq) → Array¶

Mercier stability criterion profile DMerc(s) (full mesh).

Positive interior values indicate Mercier stability. Evaluated through the parity-proven wout engine (vmex.core.nyquist.mercier_and_jxb() via wout_from_state()) — host NumPy, so this objective is finite-difference-only (not jit/AD transparent; the first two surfaces and the edge carry the usual near-axis noise, so practical targets should penalize e.g. min(DMerc[2:-1], 0)). Accepts an Equilibrium or any wout-like object.

vmex.core.optimize.l_grad_b(eq, *, s_index: int = -1, ntheta: int = 24, nphi: int = 24) → Any¶

Magnetic-gradient scale length min L_grad_B on one half-mesh surface.

L_grad_B = |B| sqrt(2 / (grad B : grad B)) with grad B : grad B the squared Frobenius norm of the Cartesian field-gradient tensor — the Kappel/Landreman coil-complexity / compactness proxy, and the L_grad_B diagnostic of the legacy QI scripts (vmex.quasi_isodynamic.objectives.lgradb_from_state). Here it is evaluated from the wout tables of the converged state: B^u/B^v from the half-mesh bsupumnc/bsupvmnc Nyquist spectra, the coordinate basis vectors and their derivatives spectrally from rmnc/zmns, and radial derivatives from the native half/full-mesh finite differences (one-sided at the edge) — so values agree with the legacy diagnostic at discretization level, not bitwise (the pointwise math lives in vmex.core.statephysics._lgradb_grid(), shared with the traceable l_grad_b_state()).

Returns the (hard) minimum over a uniform (theta, phi) grid on the selected surface (s_index indexes the ns-long half-mesh arrays; default edge). Larger is better; a practical least-squares term is max(1/L - 1/threshold, 0). Symmetric configurations only (lasym sine partners are ignored). Accepts an Equilibrium or wout-like. This lane consumes host-NumPy wout tables (finite-difference-only); use l_grad_b_state() for jac="implicit".

vmex.core.optimize.l_grad_b_state(state: SpectralState, rt: SolverRuntime, *, s_index: int = -1, ntheta: int = 24, nphi: int = 24, softmin_k: float | None = None) → Any¶

Traceable min L_grad_B of a core state (implicit-adjoint ready).

The (state, runtime) lane of l_grad_b() (Item E): identical convention — L_grad_B = |B| sqrt(2/(grad B : grad B)) minimized over the same uniform (theta, phi) grid of one half-mesh surface — with the wout coefficient tables rebuilt traceably from the state (_lgradb_state_tables(): physical rmnc/zmns from the spectral state, the wrout.f Nyquist analysis of B^u/B^v as jnp einsums) and the same half/full-mesh radial finite-difference stencils, so the default hard minimum matches the wout lane to float round-off. Fully jnp: usable directly as a two-positional objective term under jac="implicit".

softmin_k selects the reduction: None (default) is the hard min — exact, and differentiable almost everywhere (the subgradient follows the argmin point), but its gradient jumps when the minimizing gridpoint switches. A float k [1/m] returns the smooth soft minimum -logsumexp(-k * L) / k, a lower bound on the hard minimum within log(ntheta * nphi) / k (about 6.4 / k m at the default 24x24 grid; k = 50 biases a ~1 m scale length by < 0.13 m). Optimize with the smooth form, report the hard minimum.

vmex.core.optimize.quasi_isodynamic_residual(*, bmnc_b, xm_b, xn_b, iota_b, nfp: int, weights: Iterable[float] | None = None, nphi: int = 151, nalpha: int = 31, n_bounce: int = 51, include_bounce_endpoints: bool = False, softness: float = 0.02, width_weight: float = 1.0, branch_width_weight: float = 0.5, branch_width_softness: float = 0.01, profile_weight: float = 0.1, shuffle_profile_weight: float = 1.0, shuffle_profile_softness: float = 0.02, phimin: float = 0.0) → dict[str, Any]¶

Smooth Goodman-style quasi-isodynamic residual from Boozer |B| modes.

A configuration is quasi-isodynamic when the |B| contours are poloidally closed and the trapped-particle bounce distance between the two branches of each magnetic well is independent of the field-line label alpha (omnigenity). This residual samples the normalized |B| along field lines theta = alpha + iota*phi over one field period and penalizes, per surface (weights are the legacy defaults, i.e. exactly the terms the minimal-seed QI examples used):

level-set width variance (width_weight): for each bounce level B* the smooth occupancy sigmoid((B* - bnorm)/softness) gives the fraction of the field line below B*; its variance over alpha measures misalignment of the |B| contours.
branch width variance (branch_width_weight): each field line is split at its |B| minimum, both branches are made monotone with a running maximum, and the (smooth) level-crossing distances of the two branches are summed — the trapped-well bounce width, whose variance over alpha is the classic omnigenity error.
profile consistency (profile_weight): small penalty on the variance of bnorm itself over alpha at fixed phi, which keeps degenerate QH-like candidates from gaming the width terms.
branch-shuffle profile (shuffle_profile_weight): the “squash and shuffle” comparison — each well’s branch crossings are shifted so every field line has the mean bounce width, the shuffled well is reinterpolated onto the original grid and compared pointwise to the original bnorm (the closest smooth analogue of Goodman et al.’s construction of the nearest omnigenous field).

Ported from the legacy vmex.quasi_isodynamic.objectives. quasi_isodynamic_residual_from_boozer_modes with the unused aligned_profile_* / weighted_shuffle_* / shuffle_profile_nphi_out machinery removed (they defaulted to off in the QI examples). xn_b uses physical toroidal mode numbers (booz_xform convention). Returns residuals1d (least-squares vector) and total (its squared norm).

vmex.core.optimize.boozer_modes_from_wout(wout, *, surfaces, mboz: int = 18, nboz: int = 18, jit: bool = False) → dict[str, Any]¶

Boozer |B| spectrum of selected surfaces via booz_xform_jax.

wout is a WoutData (or any wout-like object accepted by Booz_xform.read_wout_data); surfaces are normalized-flux values matched to the nearest half-mesh surfaces. Returns {bmnc_b, xm_b, xn_b, iota_b, nfp, s_b} with bmnc_b shaped (nsurf, nmodes) — the inputs of quasi_isodynamic_residual().

booz_xform_jax is an optional dependency (soft import).

vmex.core.optimize.quasi_isodynamic_residual_from_wout(wout, *, surfaces, mboz: int = 18, nboz: int = 18, jit_booz: bool = False, **qi_kwargs) → dict[str, Any]¶

QI residual of a converged equilibrium: wout -> Boozer -> residual.

Convenience composition of boozer_modes_from_wout() and quasi_isodynamic_residual(); qi_kwargs are the residual’s sampling/weight knobs. Accepts a Equilibrium too, so it can be used directly as a least_squares() objective term via lambda eq: quasi_isodynamic_residual_from_wout(eq, surfaces=...)["residuals1d"].

vmex.core.optimize.boundary_dof_names(inp: VmecInput, max_mode: int) → list[str]¶

Human-readable labels (“RBC(n,m)” / “ZBS(n,m)”, INDATA index order).

For lasym boundaries the non-symmetric RBS(n,m) / ZBC(n,m) families are appended (same (m, n) order as the symmetric block).

vmex.core.optimize.pack_boundary(inp: VmecInput, max_mode: int) → ndarray¶

Flat boundary-dof vector (see _dof_modes()).

Inverse of unpack_boundary(); RBC(0,0) is excluded (fixed major radius). For a stellarator-symmetric boundary the layout is [rbc..., zbs...]; for lasym the non-symmetric families are appended as [rbc..., zbs..., rbs..., zbc...] (four families — the same m = 0 / RBC(0,0) fixing convention applies to every family, so the rigid vertical shift ZBC(0,0) and the identically-zero RBS(0,0) are excluded too).

vmex.core.optimize.unpack_boundary(inp: VmecInput, x, max_mode: int) → VmecInput¶

New VmecInput with the boundary dofs x applied.

Handles both the 2-family symmetric layout and the 4-family lasym layout (see pack_boundary()).

vmex.core.optimize.least_squares(objective_terms: Sequence[tuple[Callable, float, float]], inp: VmecInput, *, max_mode: int | Sequence[int] = 1, x0: ndarray | None = None, current_dofs: int | None = None, jac: str | None = None, jac_chunk_size: int | str | None = 'auto', jac_solver: str = 'block', recycle: bool = False, hot_restart: bool = True, warm_start: str | None = 'perturbation', use_ess: bool = False, ess_alpha: float = 1.2, device: Any = None, solve_kwargs: dict | None = None, verbose: int = 0, **scipy_kwargs)¶

Boundary-shape least squares: simsopt’s least_squares_serial_solve.

objective_terms is a list of (fun, target, weight): each fun maps a converged Equilibrium (or, for two-positional-argument callables, its (state, runtime) pair) to a scalar or residual vector, and contributes weight * (fun(eq) - target) rows to the stacked residual, i.e. cost = 1/2 sum_i w_i^2 (f_i - t_i)^2 (scipy’s 1/2 convention). The decision variables are the boundary Fourier coefficients up to max_mode (pack_boundary() — RBC(0,0) fixed). Trial boundaries whose solve fails return a large finite residual so the trust region backs off instead of crashing. Objective swaps between stages of a staged campaign (e.g. the QP-basin-then-QI route) are just two calls with different objective_terms, the second seeded with the first call’s result.input.

current_dofs = k (plan R26.g, spec section 6.4) additionally frees the equilibrium current profile: the first k AC power-series coefficients (VMEC pcurr_type="power_series", i.e. I'(s)) plus CURTOR, appended to the dof vector as [..boundary.., ac_0/ac_scale, ..., ac_{k-1}/ac_scale, curtor/1e6] (ac_scale = max(|curtor|, 1) frozen from the seed) so the trust region sees O(1) numbers. Requires ncurr = 1 and an AC-parameterized pcurr_type (splines are rejected — re-fit them to a power series first, e.g. via vmex.core.bootstrap.self_consistent_bootstrap()). Both gradient modes support it (finite differences re-solve per current dof; jac="implicit" adds k + 1 one-hot tangent rows through ImplicitParams.ac/curtor, which runtime_from_params already traces). Note that VMEC normalizes the AC profile by its own edge integral (only the shape of I' matters; CURTOR sets the amplitude), so the overall-AC-scale direction is objective-neutral; the trust-region solvers handle the resulting Jacobian null direction. This is the dof set of the self-consistent-bootstrap objective (vmex.core.bootstrap.RedlBootstrapMismatch).

max_mode may be a single int or an increasing schedule (e.g. (1, 2, 3)): each continuation stage optimizes the enlarged dof set starting from the previous stage’s boundary (higher harmonics enter at their current — typically zero — values). Repeated trial solves are cheap by construction: runtimes with the same Resolution are structural pytrees, so the solver reuses one XLA executable across all boundary trials (vmex.core.solver Phase-2 cache; only the first solve of a stage compiles). device is forwarded to the solver (vmex.core.device policy applies when None).

use_ess enables Exponential Spectral Scaling of the trust region (_ess_scale(), ess_alpha), the legacy use_ess option.

jac=None (default) uses scipy "2-point" finite differences. jac="implicit" uses the Phase-6 implicit-gradient path (vmex.core.implicit): one hot-restarted forward solve per trial boundary, and the exact residual Jacobian by forward implicit differentiation — one preconditioned GMRES per boundary dof instead of one full equilibrium solve per dof. Requirements (see the module docstring): every term traceable in (state, runtime) (vector terms expose residuals_state; wout-engine terms like d_merc() / l_grad_b() / the Boozer QI residual need jac=None — for L_grad_B use the traceable l_grad_b_state() instead). Both stellarator-symmetric and lasym boundaries are supported: the implicit parameter map packs the four RBC/ZBS/RBS/ZBC families and a traceable readin.f delta rotation. jac_chunk_size (R17.1 memory knob, jac="implicit" only) chunks the per-dof Jacobian columns via solvax.chunk_map(): "auto" (default) lets solvax.auto_chunk_size() pick a memory-bounded width (the largest block that fits the device budget on GPU, a sqrt-balanced width on CPU) so peak Jacobian memory is m0 + m1*chunk instead of scaling with the full dof count; an int fixes that many boundary dofs at a time; and None forces one wide vmap over all dofs (the pre-R17.1 behavior, fastest but peak memory O(dofs)). The column blocks are mathematically independent, so the assembled Jacobian is identical to float64 round-off (~1e-15) across chunk sizes. It is inert for jac=None (scipy computes the finite-difference Jacobian itself).

jac_solver (plan R25.2, jac="implicit" only) selects the linear solver behind the per-dof implicit-Jacobian columns. "block" (default) amortizes one block-tridiagonal factorization of the raw force Jacobian — whose radial coupling is exactly nearest-neighbor, so ns dense (3*mn, 3*mn) blocks assembled by 3-colored jax.jvp probes capture it completely at a cost independent of the dof count — then backsolves every dof right-hand side directly (solvax.block_thomas_factor() / solvax.block_thomas_solve()) and certifies each column with a short warm-started GMRES pass against the preconditioned system (same adjoint_tol norm as the default path; columns already at tolerance cost one matvec). "gmres" is the pre-R25.2 path: one independent preconditioned GMRES per dof column. Both produce the same Jacobian to solver tolerance; "gmres" is the fallback if the block path misbehaves on an exotic configuration. Inert for jac=None; recycle=True takes precedence.

recycle (plan R25.3, jac="implicit" only) carries a GCROT deflation pair across the per-dof implicit-Jacobian solves — a lax.scan over dof chunks (vmapped within a chunk) threads the solvax.gcrot() recycle space (C, U) between chunks and, via a Python-side holder, between successive trust-region Jacobian evaluations. Recycled solves keep the exact adjoint_tol / adjoint_maxiter budget of the default path. Default False: measured on the nfp2 minimal-seed max_mode-2 operator (2026-07-11), the solvax v0.1 recycle space (FIFO cycle corrections, not the harmonic Ritz vectors of GCRO-DR) slows warm-started columns — e.g. 140 (cold GMRES) -> 236/347/479 iterations at k = 2/5/10 — so columns that then exhaust adjoint_maxiter return larger residuals. Enable only after benchmarking per-column iteration counts on your operator. recycle=False uses the independent per-column solvax.gmres() path (identical columns across chunk sizes to float64 round-off). Inert for jac=None.

hot_restart seeds each trial solve from the previous converged state (both modes; in implicit mode via the per-config host-solve cache).

warm_start (plan R25.4, jac="implicit" only) refines what that seed is. "perturbation" (default) seeds each trial solve with the DESC-style first-order prediction x_ref + sum_j (dx)_j dz_j (arXiv:2203.15927 eq.perturb before eq.solve): the per-dof state responses dz_j = -(dF/dz)^{-1} dF/dp t_j are exactly the columns the implicit Jacobian already solves, so the linearization is stashed at each jac(x_ref) call for free and evaluated per trial at the cost of one small tensor contraction. "state" is the plain hot restart (seed = last converged state); None disables warm starting entirely (every trial re-solves from the interior guess). All three converge to the same fixed points — only the inner iteration count changes — and anything missing or mismatched (first call, stage change, failed seed) falls back silently down the ladder perturbation -> state -> cold. recycle=True bypasses the perturbation stash (its Jacobian variant carries the GCROT pair instead); hot_restart=False forces warm_start=None. Inert for jac=None. Measured (2026-07-12, M-series CPU, nfp2 minimal seed, max_mode=2, QS+aspect, max_nfev=20): total forward-solve iterations 23685 ("state") -> 6364 ("perturbation"), 3.7x fewer, wall 156 s -> 145 s (this deck’s wall is Jacobian-dominated at ns=35; the solve-phase win grows with ns), and the seed additionally rescued a trial whose plain hot restart failed with a Jacobian-sign error.

Remaining keywords go to scipy.optimize.least_squares() (e.g. max_nfev, ftol, xtol, diff_step).

Returns the scipy OptimizeResult of the final stage with extra attributes: input (optimized VmecInput), equilibrium (last successfully solved Equilibrium), stage_results (per-max_mode results for schedules) and, in implicit mode, solve_stats ({"solves", "iterations"} totals of the stage’s host forward solves — the R25.4 warm-start instrumentation).

Concurrent ensembles of independent equilibrium solves on CPU (Item G).

A parameter scan or an ensemble optimization solves N independent equilibria (different boundaries / phiedge / profiles). Each forward solve runs on the host behind jax.pure_callback() (vmex.core.implicit) or directly through vmex.core.solver.solve() / vmex.core.multigrid.solve_multigrid(), and — crucially — releases the Python GIL while XLA executes the compiled iteration lanes. A plain concurrent.futures.ThreadPoolExecutor over those independent solves therefore overlaps their XLA execution and gives real wall-clock speedup, while every result stays byte-identical to solving that input alone (the solves share no mutable state).

Measured strong scaling (this box: 10 logical CPUs, nfp2_QA phiedge scan, 8 balanced solves ~0.68 s each, best-of-3):

workers   wall   speedup   efficiency
serial   5.46 s   1.00x       100 %
2        3.05 s   1.79x        89 %
4        2.15 s   2.54x        63 %
8        1.66 s   3.29x        41 %

The scaling is deliberately sub-linear: XLA already multithreads within one solve, so as the worker count approaches the core count the per-solve XLA threads contend — the ensemble speedup and the intra-solve speedup draw from the same cores. See Parallelization for the full mechanism study (why threading beats pmap across forced host devices and vmap over the callback here), the honest limits (Amdahl on imbalanced heterogeneous ensembles; the launch-bound implicit adjoint overlaps far less than the forward solve), and the multi-GPU design sketch.

This module is a thin, additive concurrency layer: it changes nothing in the single-solve path (which stays byte-identical) and imposes no new dependency.

vmex.core.parallel.default_workers(n_items: int, workers: int | None = None) → int¶

Resolve the worker count for an n_items ensemble.

None (default) picks min(n_items, os.cpu_count()) — enough threads to cover the ensemble without oversubscribing the cores that each solve’s XLA threads already use. An explicit workers is honoured but clamped to [1, n_items] (more threads than items cannot help, and 0/negative is meaningless).

vmex.core.parallel.map_ensemble(fn: Callable[[_T], _R], items: Iterable[_T], *, workers: int | None = None, return_exceptions: bool = False) → list[_R]¶

Apply fn to each of items concurrently on CPU; keep input order.

The general primitive behind solve_ensemble(). fn must be an independent per-item computation (e.g. a full vj.solve / implicit.run / jax.value_and_grad over one input) that shares no mutable state with the others — which every vmex forward solve is, since each builds its own runtime and the compiled-executable cache is thread-safe. Under those conditions the results are byte-identical to a serial [fn(x) for x in items] (the concurrency only overlaps the GIL-releasing XLA execution windows).

workers — see default_workers(). With workers=1 the pool runs sequentially (a clean serial baseline for scaling measurements).

return_exceptions=False (default) re-raises the first item’s exception (preserving vmex’s typed VmecError taxonomy), exactly as a serial loop would. return_exceptions=True instead places the caught exception object in that slot so one failed ensemble member does not abort the batch (useful for optimization ensembles / robustness scans).

vmex.core.parallel.solve_ensemble(inputs: Sequence[Any], *, workers: int | None = None, multigrid: bool = True, return_exceptions: bool = False, **solve_kwargs: Any) → list[Any]¶

Solve N independent VmecInput concurrently.

Threads vmex.core.multigrid.solve_multigrid() (multigrid=True, default — runs each input’s NS_ARRAY ladder) or vmex.core.solver.solve() (multigrid=False, single grid) over the ensemble on CPU, returning the list of SolveResult in input order. Each result is byte-identical to solving that input by itself (verified in tests/test_parallel.py): the helper only overlaps the solves’ XLA execution — it does not touch the numerics, the convergence path, or the default single-solve code path.

Extra **solve_kwargs (e.g. verbose, ftol, initial_state) are forwarded unchanged to every solve. workers and return_exceptions behave as in map_ensemble().

Best speedup comes from a balanced ensemble — a parameter scan at fixed resolution, where the members share a compiled executable and take a similar iteration count. A heterogeneous ensemble is limited by its slowest member (Amdahl); see Parallelization.

Physics objectives¶

The objective catalog with usage snippets is Objectives library.

Traceable omnigenity/QI objective (R26h.h2).

A quasi-isodynamic (QI) omnigenity residual evaluated as a pure, traceable function of a converged (SpectralState, SolverRuntime) pair, so it composes with the implicit-gradient least-squares lane (vmex.core.optimize.least_squares(..., jac="implicit")) exactly like QuasisymmetryRatioResidual — no wout tables, no host booz_xform round-trip.

Two pieces:

boozer_bmnc_state() — a minimal traceable Boozer |B| transform. The classic BOOZ_XFORM construction (Hirshman’s Fortran code; equation numbers below follow Landreman’s booz_xform documentation) is evaluated in pure jax.numpy from the solver’s internal half-mesh field tables: the periodic part of the Boozer generating potential w is integrated spectrally from the covariant field components (dw/dtheta = B_theta, dw/dzeta = B_zeta), then nu = (w - I*lambda) / (G + iota*I) gives the Boozer angles theta_B = theta + lambda + iota*nu, zeta_B = zeta + nu, and the Boozer |B| harmonics come from the angle-transform quadrature bmnc_b = <|B| cos(m theta_B - n zeta_B) * d(theta_B, zeta_B)/d(theta, zeta)> — the same equations booz_xform solves, here end-to-end differentiable.
omnigenity_residual() / QIResidual — a smooth omnigenity distance for poloidally-closed-contour (M = 0, N = 1) omnigenity. The formulation is the constructed-QI-target distance of Goodman et al., “Constructing precisely quasi-isodynamic magnetic fields”, J. Plasma Phys. 89, 905890504 (2023), arXiv:2211.09829, distilled to its level-set form: on each surface, |B| is sampled along Boozer field lines theta_B = alpha + iota * phi_B over one field period and the residual stacks, per surface,
- bounce-distance uniformity (well_weight): for every trapping level B*, the bounce distance delta(alpha, B*) between the two monotone branches of the magnetic well (smooth occupancy integrals of the running-maximum branch envelopes) minus its field-line average — the Cary–Shasharina omnigenity condition (Cary & Shasharina, PRL 78, 674 (1997)) that Goodman’s “shuffle” step enforces;
- extremum alignment (extremum_weight): the per-field-line B_min/B_max minus their field-line averages — poloidal closure of the extremal |B| contours (Goodman’s “align the maxima” step; also the flat-B_max condition of Dudt et al., J. Plasma Phys. 90, 905900120 (2024), arXiv:2305.08026);
- single-well monotonicity (squash_weight): the pointwise distance between |B| and its monotone branch envelopes — Goodman’s “squash” distance, penalizing side extrema (multiple wells per period).
Each piece is an exact zero of an exactly QI field, every operation is smooth or piecewise-smooth (sigmoid occupancies, running maxima), and the full pipeline is jit/grad/jvp-transparent for the implicit lane.

Scope notes¶

Stellarator-symmetric states only (lasym = False), like the other traceable objectives.
Requested surfaces are snapped to the nearest half-mesh surface (the Boozer transform is a per-surface construction — same convention as vmex.core.optimize.boozer_modes_from_wout()), not interpolated.
The wout-engine analogue for cross-checks is vmex.core.optimize.quasi_isodynamic_residual_from_wout() (host booz_xform_jax, finite-difference-only).

vmex.core.omnigenity.boozer_bmnc_state(state: SpectralState, rt: SolverRuntime, *, surfaces, mboz: int = 16, nboz: int = 16, oversample: int = 2) → dict[str, Any]¶

Boozer |B| cosine spectrum of selected surfaces, fully traceable.

The jnp analogue of vmex.core.optimize.boozer_modes_from_wout() (booz_xform algorithm, module docstring) evaluated from the solver’s internal tables: |B|, the covariant components B_theta/B_zeta and the Boozer profile averages I/G live on the half-mesh (theta, zeta) grid (mirrored to the full theta circle), lambda comes from the state’s spectral coefficients with the exact wout half-mesh rescale, and the generating potential w is integrated spectrally (FFT) — m != 0 modes from B_theta, m = 0 modes from B_zeta, the booz_xform mode split.

surfaces are normalized-flux values snapped to the nearest half-mesh surfaces (one Boozer construction per requested value, duplicates kept so outputs align with surfaces). oversample refines the quadrature grid by trigonometric (FFT zero-pad) interpolation before the Boozer angle transform, reducing the aliasing of cos(m theta_B - n zeta_B) products; mboz/nboz are capped at the fine grid’s Nyquist.

Returns {bmnc_b (nsurf, nmodes), xm_b, xn_b (physical), iota_b, nfp, s_b} — bmnc_b/xm_b/xn_b/iota_b/nfp are the spectrum inputs of omnigenity_residual() and of vmex.core.optimize.quasi_isodynamic_residual().

vmex.core.omnigenity.omnigenity_residual(*, bmnc_b, xm_b, xn_b, iota_b, nfp: int, weights: Iterable[float] | None = None, nphi: int = 97, nalpha: int = 25, n_levels: int = 16, softness: float = 0.02, well_weight: float = 1.0, extremum_weight: float = 1.0, squash_weight: float = 1.0) → dict[str, Any]¶

Smooth constructed-QI-target omnigenity residual (module docstring).

|B| is synthesized from Boozer harmonics along field lines theta_B = alpha + iota * phi_B on nalpha labels over one field period (nphi periodic points), normalized per surface to [0, 1] by the surface extrema. Each field line is split at its minimum into two monotone branch envelopes (periodic running maxima); the residual stacks the three Goodman-construction distances:

well: bounce distance delta(alpha, B*) = d_left + d_right (sigmoid occupancy integrals of the branch envelopes at n_levels trapping levels, in field-period fraction units) minus its alpha-average — Cary–Shasharina bounce-distance omnigenity;
extremum: per-line min/max of |B| minus their alpha-averages — poloidally closed extremal contours;
squash: pointwise envelope - |B| monotonicity defect — one magnetic well per field period.

All three vanish on an exactly QI field. softness is the sigmoid level width in normalized |B| units. Returns residuals1d (flat least-squares vector), total = sum(residuals1d**2) and diagnostics.

class vmex.core.omnigenity.QIResidual(surfaces, *, weights: Iterable[float] | None = None, mboz: int = 16, nboz: int = 16, oversample: int = 2, nphi: int = 97, nalpha: int = 25, n_levels: int = 16, softness: float = 0.02, well_weight: float = 1.0, extremum_weight: float = 1.0, squash_weight: float = 1.0)¶

Traceable quasi-isodynamic omnigenity residual (module docstring).

Composition of boozer_bmnc_state() (traceable Boozer |B| spectrum on the requested surfaces) and omnigenity_residual() (smooth Goodman constructed-QI-target distance). The interface mirrors QuasisymmetryRatioResidual: the instance is a least_squares() objective term for both gradient modes — jac=None calls J() on the converged Equilibrium, jac="implicit" picks up the traceable residuals_state() vector (full pointwise Gauss-Newton geometry, exact implicit gradients).

Example:

qi = QIResidual(np.linspace(0.25, 0.9, 4))
result = least_squares([(qi, 0.0, 10.0), ...], inp, jac="implicit")

compute_state(state: SpectralState, rt: SolverRuntime) → dict[str, Any]¶: Full diagnostics dict (Boozer spectrum + residual pieces).

residuals_state(state: SpectralState, rt: SolverRuntime) → Array¶: Traceable flat residual vector with sum(r**2) = total_state.

total_state(state: SpectralState, rt: SolverRuntime) → Any¶: Traceable scalar omnigenity objective sum(residuals**2).

J(eq) → Array¶: Objective-term entry point for least_squares (residual vector).

residuals(eq) → Array¶: Alias of J() (simsopt-style vocabulary).

total(eq) → Any¶: Scalar omnigenity objective of a converged equilibrium.

Differentiable Redl (2021) bootstrap current (R26.g, steps 1-2).

Implements the pure-formula layer and the two geometry lanes of notes_r26g_redl_spec.md:

compute_trapped_fraction() — effective trapped fraction, epsilon and the flux-surface averages <B^2>, <1/B> from |B|/sqrt(g) on an angular grid. Differentiable rewrite of simsopt simsopt.mhd.bootstrap.compute_trapped_fraction (fixed-order Gauss-Legendre lambda quadrature instead of adaptive scipy.integrate.quad; plain grid max/min instead of spline-refined extrema; double-where guard on the sqrt(1 - lambda B) near-singularity so reverse-mode AD stays finite).
j_dot_B_redl() — the Redl et al., Phys. Plasmas 28, 022502 (2021) <J.B> formula (eqs. 10-16, 19-21 with the Sauter 18b-18e collisionalities), transcribed verbatim from simsopt j_dot_B_Redl including the quasisymmetry isomorphism as simsopt applies it: iota -> iota - nfp*helicity_n everywhere iota appears, G unshifted (spec section 3 note).
KineticProfiles / profile_value_and_dds() — polynomial ne/Te/Ti/Zeff profiles in s (lowest order first, simsopt ProfilePolynomial convention), value + analytic d/ds via Horner.
redl_geometry_from_wout() — parity lane mirroring simsopt RedlGeomVmec.__call__: half-mesh linear interpolation of iotas/bvco/buco/gmnc/bmnc onto the requested surfaces, cosine synthesis of |B|/sqrt(g) on a uniform (theta, phi) grid, then compute_trapped_fraction().
redl_geometry_from_state() — traceable lane on (SpectralState, SolverRuntime): |B| and sqrt(g) from the solver’s half-mesh internal grid (mirrored from the reduced [0, pi] theta grid exactly as QuasisymmetryRatioResidual._pointwise_state), iota/G/I from _iotas_half/surface_currents.

Steps 3-4 of the spec (this module, second landing):

vmec_j_dot_B() / vmec_j_dot_B_from_wout() — the equilibrium <J.B> via the MHD identity (spec section 6.2, validated against the Zenodo convertSfincsToVmecCurrentProfile script and the wout jdotb):

<J.B>(s) = [<B^2> dI/ds + mu0 I dp/ds] / (2 pi psi_a), I(s) = signgs*(2 pi/mu0)*buco(s) [A], psi_a = phi(1)/(2 pi) [Wb/rad]

(the signgs matches VMEC’s ctor = signgs*(2 pi/mu0)*buco(ns) convention, so the identity reproduces the sign of the wout jdotb).
RedlBootstrapMismatch — the paper/simsopt f_boot normalized residual R_j = (Jv_j - Jr_j)/sqrt(sum_k (Jv_k + Jr_k)^2) comparing <J.B>_vmec against j_dot_B_redl(), with the same wout / traceable-state dual-lane shape as QuasisymmetryRatioResidual (so it composes with least_squares at jac=None and jac="implicit").
self_consistent_bootstrap() — fixed-boundary Picard iteration AC/CURTOR <- current profile implied by the Redl <J.B> (the Zenodo script’s “smooth method”: solve [<B^2> d/ds + mu0 dp/ds] I = 2 pi psi_a J_Redl with I(0) = 0, then re-fit the I'(s) power series).

Units follow simsopt/the spec: ne [1/m^3], Te/Ti [eV], G/I/R [T*m], psi_edge [Wb/rad], output <J.B> [A*T/m^2].

Note (spec section 6.1b): the spec’s stated decision was to hoist _field_chain/_iotas_half/_half_grid/_interp_half_grid into a shared module; since R26a they live in vmex.core.statephysics (optimize re-exports them for backward compatibility).

vmex.core.bootstrap.ELEMENTARY_CHARGE = 1.602176634e-19¶: CODATA elementary charge [C]; converts eV -> J in the <J.B> assembly.

class vmex.core.bootstrap.KineticProfiles(ne_coeffs: Any, Te_coeffs: Any, Ti_coeffs: Any, Zeff_coeffs: Any = 1.0)¶

Prescribed kinetic profiles for the Redl formula (frozen pytree).

Polynomial coefficients in s (lowest order first). These are not VMEC inputs; they parameterize the bootstrap objective (spec section 6.6). Units: ne_coeffs [1/m^3], Te_coeffs/Ti_coeffs [eV], Zeff_coeffs dimensionless (default: constant 1, hydrogen).

Example (paper profiles): ne = n0*(1 - s^5), Te = Ti = T0*(1 - s):

KineticProfiles(ne_coeffs=4.13e20 * np.array([1, 0, 0, 0, 0, -1]),
                Te_coeffs=12.0e3 * np.array([1, -1]),
                Ti_coeffs=12.0e3 * np.array([1, -1]))

vmex.core.bootstrap.profile_value_and_dds(coeffs: Any, s: Any) → tuple[Any, Any]¶

Polynomial profile value and analytic d/ds at s (traceable).

coeffs are polynomial coefficients in s, lowest order first (simsopt ProfilePolynomial): p(s) = sum_k coeffs[k] * s**k. Evaluated with the paired Horner recurrence, so the derivative is exact and both outputs are differentiable in coeffs and s.

vmex.core.bootstrap.compute_trapped_fraction(modB: Any, sqrtg: Any, *, n_lambda: int = 64)¶

Effective trapped fraction and flux-surface averages per surface.

f_t = 1 - (3/4) <B^2> \int_0^{1/Bmax} lambda dlambda / <sqrt(1 - lambda B)> with <.> the flux-surface average (weight |sqrt(g)|; only ratios of averages enter, so the uniform Jacobian sign cancels).

Parameters:

modB – (nsurf, ntheta, nzeta) array of \(|B|\) on a uniform angular grid (leading surface axis — note simsopt uses trailing).
sqrtg – same shape, the Jacobian \(\sqrt{g}\) on the grid.
n_lambda – fixed Gauss-Legendre quadrature order for the lambda integral on [0, 1/Bmax] (replaces simsopt’s adaptive quad; the nodes exclude the lambda = 1/Bmax endpoint).

Returns:

(Bmin, Bmax, epsilon, fsa_B2, fsa_1overB, f_t) — 1D arrays of length nsurf (same tuple as simsopt). Bmin/Bmax are hard grid extrema: piecewise-smooth gradients, adequate for trust-region least squares (same stance as vmex.core.optimize.mirror_ratio()).

class vmex.core.bootstrap.RedlGeometry(surfaces: Any, iota: Any, G: Any, I: Any, R: Any, epsilon: Any, f_t: Any, fsa_B2: Any, fsa_1overB: Any, Bmin: Any, Bmax: Any, psi_edge: Any, nfp: int = 1)¶

Per-surface geometry inputs of the Redl formula (frozen pytree).

All arrays are 1D over surfaces (values of normalized toroidal flux). G/I are the Boozer covariant averages (wout bvco/buco, [T*m]); R = (G + iota*I) * <1/B> is the effective major radius for the Sauter collisionality; psi_edge = -phi(s=1)/(2 pi) [Wb/rad] (wout sign convention). nfp is static metadata.

vmex.core.bootstrap.redl_geometry_from_wout(wout, surfaces, *, ntheta: int = 64, nphi: int = 65, n_lambda: int = 64) → RedlGeometry¶

Redl geometry inputs from a wout dataset (simsopt RedlGeomVmec lane).

wout may be a WoutData, a path to a wout_*.nc file, or anything exposing a .wout attribute (e.g. Equilibrium). Mirrors simsopt RedlGeomVmec.__call__: linear interpolation of the half-mesh tables iotas/bvco/buco/gmnc/bmnc (axis slot dropped) onto surfaces, cosine synthesis of |B|/sqrt(g) from the Nyquist mode tables on theta in [0, 2 pi) x phi in [0, 2 pi/nfp) grids, then compute_trapped_fraction(). jnp throughout; this is the validation/finite-difference lane, not the implicit path.

vmex.core.bootstrap.redl_geometry_from_state(state: SpectralState, rt: SolverRuntime, *, surfaces=None, n_lambda: int = 64) → RedlGeometry¶

Redl geometry inputs as a traceable function of (state, runtime).

|B| and sqrt(g) are taken on the solver’s half-mesh internal grid (see _half_mesh_fields()); psi_edge = -signgs*hs*sum(phipf[1:]). Per-surface scalars and the angular fields are linearly interpolated from the half mesh onto surfaces before compute_trapped_fraction().

Default surfaces: linspace(0.05, 0.95, 16) — interior only; do not sample s -> 1 where Te, Ti -> 0 blows up the collisionality (spec section 6.1b).

vmex.core.bootstrap.j_dot_B_redl(profiles: KineticProfiles, geom: RedlGeometry, helicity_n: int) → tuple[Any, dict[str, Any]]¶

Bootstrap <J.B> [A*T/m^2] from the Redl (2021) formulae (pure jnp).

Equation-by-equation transcription of simsopt j_dot_B_Redl (which generated the arXiv:2205.02914 results): Sauter eqs. (18b)-(18e) for the Coulomb logarithms and collisionalities, Redl eqs. (10)-(16) and (19)-(21) for L31/L32/L34/alpha. The quasisymmetry isomorphism is applied exactly as simsopt does: iota -> iota - N with N = nfp*helicity_n in the collisionality geometry factor and in the 1/(psi_edge*(iota - N)) prefactor; G (and the R provided by the geometry lanes) is used unshifted.

helicity_n is 0 for quasi-axisymmetry, +/-1 for quasi-helical symmetry. Te/Ti are clamped at 1 eV and ne at 1e17 1/m^3 (belt-and-suspenders against sampling s -> 1 where the profiles vanish and the collisionality blows up; spec section 6.1b).

Returns (jdotB, details) with details a dict of every intermediate quantity (simsopt names).

vmex.core.bootstrap.vmec_j_dot_B(state: SpectralState, rt: SolverRuntime, *, surfaces=None) → Any¶

Traceable equilibrium <J.B> [A*T/m^2] via the MHD identity.

For a VMEC equilibrium the flux-surface-averaged parallel current obeys (spec section 6.2; validated in the Zenodo convertSfincsToVmecCurrentProfile script against VMEC’s own jdotb)

<J.B>(s) = [<B^2> dI/ds + mu0 I dp/ds] / (2 pi psi_a)

with I(s) = signgs*(2 pi/mu0)*buco(s) [A] (VMEC’s ctor = signgs*(2 pi/mu0)*buco(ns) sign convention), p in Pa and psi_a = phi(1)/(2 pi) [Wb/rad]. In core internal units (fields.pressure is mu0*Pa, buco [T*m]) this reduces to signgs*(<B^2> buco' + buco p_int') / (mu0 psi_a). Radial derivatives are half-mesh central differences (one-sided at the ends), then the half-mesh profile is linearly interpolated onto surfaces (default: linspace(0.05, 0.95, 16), matching redl_geometry_from_state()).

Agrees with the wout-engine jdotb (jxbforce.f) to a few 1e-4 relative on the Zenodo finite-beta optima — no bsubs port needed.

vmex.core.bootstrap.vmec_j_dot_B_from_wout(wout, surfaces, *, geom: RedlGeometry | None = None, ntheta: int = 64, nphi: int = 65) → Any¶

The section-6.2 identity evaluated from wout tables (parity lane).

Same formula as vmec_j_dot_B() with buco/pres/phi read from the wout dataset and <B^2> synthesized from bmnc/gmnc at the requested surfaces; pass geom (a RedlGeometry from redl_geometry_from_wout() at the same surfaces) to reuse its fsa_B2 instead. This is the validation lane — the wout jdotb itself (jxbforce.f) is what RedlBootstrapMismatch consumes.

class vmex.core.bootstrap.RedlBootstrapMismatch(profiles: KineticProfiles, helicity_n: int, surfaces=None, *, ntheta: int = 64, nphi: int = 65, n_lambda: int = 64)¶

Normalized <J.B>_vmec vs <J.B>_Redl mismatch (paper f_boot).

Residual vector (simsopt VmecRedlBootstrapMismatch.residuals verbatim):

R_j = (Jv(s_j) - Jr(s_j)) / sqrt(sum_k (Jv(s_k) + Jr(s_k))**2)

over the geometry surfaces, so sum(R**2) is the paper’s

f_boot = int ds [<J.B>_vmec - <J.B>_Redl]^2
/ int ds [<J.B>_vmec + <J.B>_Redl]^2

(bounded by 1; the denominator depends on the dofs — the simsopt/paper self-normalizing convention, kept as is). Dual-lane shape mirroring QuasisymmetryRatioResidual:

J()/residuals() (wout lane, FD-only): Jr from redl_geometry_from_wout(), Jv from the wout jdotb (jxbforce.f) interpolated onto surfaces — simsopt parity.
residuals_state() (traceable lane): one _half_mesh_fields() pass feeds both the Redl geometry and the section-6.2 identity Jv, so the term composes with least_squares(jac="implicit") (profiles fixed, geometry traced).

helicity_n is 0 for QA, -1/+1 for QH (units of nfp, simsopt convention). Default surfaces: linspace(0.05, 0.95, 16) (interior only — never sample s -> 1; spec section 6.1b). Usage as an objective term (target 0; the normalization already scales it):

terms = [(qs.residuals_state, 0.0, 1.0),
         (boot.residuals_state, 0.0, w_boot), ...]

residuals(wout) → Array¶: Residual vector from a wout-like object / path / Equilibrium.

profile(wout) → Array¶: Per-surface squared residuals (sum = total).

total(wout) → Any¶: Scalar f_boot = sum(residuals**2).

J(eq) → Array¶: Objective-term entry point for least_squares().

residuals_state(state: SpectralState, rt: SolverRuntime) → Array¶

Traceable residual vector (implicit lane; sum(r**2) = total_state).

Same R_j with Jr from redl_geometry_from_state() and Jv from vmec_j_dot_B() — one shared field-chain pass. Agrees with residuals() at discretization level (solver internal grid + the identity Jv vs the 64x65 wout synthesis + jxbforce jdotb).

profile_state(state: SpectralState, rt: SolverRuntime) → Any¶: Traceable per-surface squared residuals.

total_state(state: SpectralState, rt: SolverRuntime) → Any¶: Traceable scalar f_boot.

class vmex.core.bootstrap.BootstrapPicardResult(input: VmecInput, equilibrium: Any, converged: bool, iterations: int, history: tuple)¶

Outcome of self_consistent_bootstrap().

input carries the final AC/CURTOR (ncurr=1, pcurr_type="power_series"); equilibrium is the last solved Equilibrium; history is one dict per iteration (curtor, delta, f_boot). delta is the self-consistency error max|I'_Redl - I'_applied| / max|I'_Redl| between the Redl-implied and the currently applied current profile.

vmex.core.bootstrap.self_consistent_bootstrap(inp: VmecInput, profiles: KineticProfiles, helicity_n: int, *, n_iter: int = 10, tol: float = 0.001, relax: float = 1.0, degree: int = 12, s_eval=None, solve_kwargs: dict | None = None, verbose: bool = False) → BootstrapPicardResult¶

Fixed-boundary Picard iteration to a bootstrap-consistent current profile.

The Zenodo convertSfincsToVmecCurrentProfile “smooth method” with the Redl <J.B> in place of SFINCS (spec section 6.5, secondary lane). Host-side loop (no AD through it):

solve the equilibrium (hot-restarted from the previous iterate);
Jr = j_dot_B_redl(profiles, redl_geometry_from_wout(...)) on the interior s_eval grid (default linspace(0.02, 0.98, 49)), pinned to 0 at s = 0, 1 and interpolated onto the full mesh;
invert the section-6.2 identity for the enclosed current: solve [<B^2> d/ds + mu0 dp/ds] I = 2 pi psi_a Jr with I(0) = 0 (dense collocation solve, _picard_dds_matrix()), then recompute the smooth dI/ds = (2 pi psi_a Jr - mu0 I dp/ds)/<B^2>;
under-relax I' <- (1 - relax) I'_prev + relax I'_new, refit the VMEC AC power series (numpy.polynomial.polyfit, ascending — pcurr_type="power_series" parameterizes I') and set CURTOR = I(1);
stop when delta = max|I'_new - I'_applied| / max|I'_new| <= tol.

relax = 1.0 is the plain fixed point (paper expectation: <= 5 iterations at beta <= 2.5%); use relax = 0.5 when the current dominates iota (e.g. tokamaks / the beta = 5% QH), where the <J.B> ~ 1/iota ~ 1/I feedback makes the undamped map marginally stable. degree is clipped to ns - 2. The first solve uses inp’s own current settings (any pcurr_type); every refit switches to ncurr=1 + power_series.

Differentiable ideal-MHD stability objectives (R26h.h1).

Infinite-n ideal-ballooning growth rate as a pure, traceable function of a converged (SpectralState, SolverRuntime) pair — the JAX analogue of the COBRA solve (Sanchez, Hirshman, Ware, Berry & Batchelor, J. Comput. Phys. 161, 589 (2000)) in the modern differentiable formulation of Gaur et al., J. Plasma Phys. 89, 905890518 (2023) (arXiv:2302.07673):

d/dη ( g dX/dη ) + c X = λ f X , X(±η_b) = 0 , g, f > 0,

a self-adjoint second-order ODE eigenproblem along a field line, with η the straight-field-line (PEST) poloidal angle. λ > 0 is unstable; λ is the squared growth rate normalized to the Alfvén frequency, λ = (γ a_N / v_A)² with a_N the effective minor radius and v_A the Alfvén speed at the reference field B_N = 2|ψ_edge|/a_N².

The coefficient arrays follow the (COBRA-validated) VMEC-coordinate geometry conventions of simsopt’s vmec_fieldlines (Landreman) and of the adjoint ballooning solver of Gaur et al.:

g = |b·∇η| |∇α|² a_N³ B_N / B² (field-line bending),
c = -2 μ0 (dp/dρ) (b×κ·∇α)-type pressure/curvature drive, assembled from B×∇|B|·∇α plus the μ0 dp/dψ force-balance correction,
f = |∇α|² a_N B_N³ / (B² |B·∇η| B) (inertia),

with α = θ* - ι (φ - ζ0) the field-line label, ∇α carrying the secular magnetic-shear term -ι'(s) (φ - ζ0) ∇s (ballooning parameter ζ0). All geometry is evaluated spectrally from the converged state’s physical Fourier coefficients: values and angular derivatives are exact trig sums; radial dependence uses a local parabola through the three neighbouring full-mesh surfaces, and every derivative (including the second derivatives inside ∇|B|) comes from JAX automatic differentiation of the point-evaluation functions, so the whole pipeline is jit/grad-transparent.

The discretized problem is the symmetric tridiagonal generalized eigenproblem of the central-difference stencil (the same discretization as COBRA/DESC), reduced to standard form with the diagonal f^{-1/2} similarity transform and solved with a batched jnp.linalg.eigvalsh over all requested (surface, α, ζ0) field lines.

Scope notes¶

Stellarator-symmetric states only (lasym = False), matching the other traceable objectives in vmex.core.optimize.
Surfaces need ι ≠ 0 (the field-line parameterization divides by ι).
A traceable Mercier DMerc was considered for this pass and cut: the parity-proven port (vmex.core.nyquist.mercier_and_jxb()) needs the jxbforce-filtered covariant field tables and bsubs synthesis (host-side Nyquist machinery), so a faithful jnp translation is a larger follow-up. Use vmex.core.optimize.d_merc() (wout engine, finite-difference-only) and vmex.core.optimize.magnetic_well() (traceable) for Mercier-adjacent targets meanwhile.

vmex.core.stability.ballooning_lambda(state: SpectralState, rt: SolverRuntime, *, s_indices: Sequence[int] | None = None, alphas: Sequence[float] | None = None, zeta0s: Sequence[float] = (0.0,), npoints: int = 121, nturns: float = 3.0) → Array¶

Most-unstable ideal-ballooning eigenvalue per field line (traceable).

Solves the infinite-n ideal-ballooning eigenproblem (module docstring) on the requested full-mesh surfaces and field lines of a converged core state and returns the largest eigenvalue λ of each line, shaped (len(s_indices), len(alphas), len(zeta0s)). λ > 0 means ballooning-unstable with normalized squared growth rate λ = (γ a_N/v_A)²; λ < 0 is the stable (oscillating) side.

Parameters:

s_indices – Full-mesh surface indices, each in [2, ns - 2]. Default: three surfaces at ~35/60/85 % of the radius.
alphas – Field-line labels α = θ* - ι (φ - ζ0). Default: four lines uniform in [0, π] (stellarator symmetry maps α -> -α).
zeta0s – Ballooning parameters (the toroidal angle where the secular radial wavenumber vanishes). Default (0,).
npoints – Field-line grid: npoints points over θ* ∈ α ± nturns·π (COBRA-style domain; Gaur et al. use 5π, 3 turns is adequate for optimization-grade accuracy at these resolutions).
nturns – Field-line grid: npoints points over θ* ∈ α ± nturns·π (COBRA-style domain; Gaur et al. use 5π, 3 turns is adequate for optimization-grade accuracy at these resolutions).

vmex.core.stability.ballooning_growth_rate(state: SpectralState, rt: SolverRuntime, *, s_indices: Sequence[int] | None = None, alphas: Sequence[float] | None = None, zeta0s: Sequence[float] = (0.0,), npoints: int = 121, nturns: float = 3.0, reduction: str = 'softmax', temperature: float = 0.05) → Any¶

Scalar ballooning objective: (smooth) max of λ over all field lines.

Reduces ballooning_lambda() over surfaces × field lines: reduction="softmax" (default) uses the smooth upper bound T · logsumexp(λ/T) (within T·log(N) of the hard max, fully AD-friendly for the implicit-gradient lane); "max" returns the hard maximum. Drive it negative (e.g. target -0.01 in vmex.core.optimize.least_squares()) for stable-by-construction campaigns; the signature is a two-positional (state, runtime) callable, so it works with both jac=None and jac="implicit".

Turbulence-proxy optimization objectives via SPECTRAX-GK (R26h.h4).

Wires the gyrokinetic turbulence proxies of SPECTRAX-GK (uwplasma; JAX-native Hermite-Laguerre flux-tube solver) to converged (SpectralState, SolverRuntime) pairs, in two layers:

Geometry adapter — gk_fieldline_geometry() samples one field line of the converged interior solution and emits the solver-ready flux-tube geometry contract of spectraxgk.flux_tube_geometry_from_mapping (bmag, gradpar, gds2/gds21/gds22, gbdrift/gbdrift0, cvdrift/cvdrift0, bgrad, …). The arrays follow the GS2/GX normalizations of simsopt’s vmec_fieldlines (Landreman) — the exact conventions already used by the ballooning objective in vmex.core.stability, whose spectral point-evaluation machinery (_ballooning_context, _parabola, _theta_vmec_from_pest) is reused here, extended with the grad s/grad psi metric and drift projections that ballooning does not need. Everything is exact trig sums + JAX AD: the adapter is jit/grad-transparent and needs no spectraxgk import.
Objective wrappers — thin (state, runtime) callables around the proxies SPECTRAX-GK itself promotes for VMEC-side optimization (its VMECJAXTransportObjectiveConfig kinds, docs stellarator_optimization.rst):
- turbulent_growth_rate() — kind "growth": dominant linear ITG/TEM-branch growth rate gamma of the spectral gyrokinetic operator on the sampled flux tube (spectraxgk.solver_growth_rate_from_geometry; the eigenvalue carries SPECTRAX-GK’s implicit-eigenpair custom AD rule).
- quasilinear_flux_proxy() — kind "quasilinear_flux": the mixing-length quasilinear heat-flux proxy gamma * W_Q / k_perp_eff^2 built from the dominant eigenmode’s heat-flux weight and effective perpendicular wavenumber.
- nonlinear_heat_flux_proxy() — kind "nonlinear_window_heat_flux": SPECTRAX-GK’s smooth reduced nonlinear-window heat-flux surrogate (saturation-rule closure csat * W_Q * 2 gamma_+ / (1 + 2.2 k_perp_eff^2 + 0.15 gamma_+), spectraxgk.objectives.vmec_transport_tables). This is the documented proxy for the nonlinear transport window; a production nonlinear claim still requires SPECTRAX-GK’s matched long nonlinear audits, per its own docs.
- turbulence_objective_vector() — the underlying ordered SOLVER_OBJECTIVE_NAMES vector (gamma, omega, kperp_eff2, linear_heat_flux_weight, linear_particle_flux_weight, mixing_length_heat_flux_proxy).
Each wrapper is a two-positional (state, runtime) callable, so it composes with vmex.core.optimize.least_squares(). Traceability status (validated in tests/test_turbulence.py): SPECTRAX-GK is JAX-native, and turbulent_growth_rate() is fully differentiable in both AD modes (jac=None and jac="implicit" — the wrapper reduces SPECTRAX-GK’s explicit operator matrix with jnp.linalg.eigvals, which carries a JVP, where SPECTRAX-GK’s own dominant_real_eigenvalue is a reverse-only custom_vjp that the forward-mode implicit Jacobian cannot trace). The quasilinear and nonlinear-window proxies additionally weight the dominant eigenvector (heat-flux weight, kperp_eff), and JAX declines derivatives of non-symmetric eigenvectors — those two are value-level objectives: use jac=None (finite differences), exactly like the wout-engine terms (d_merc, l_grad_b) in the optimization examples.

The heavy dependency is optional: only the objective wrappers import spectraxgk (pip install spectraxgk; its solvax pin is satisfied API-wise by the in-house solvax’s gmres/tridiagonal_solve/ chunked_jacfwd). The geometry adapter works without it.

Scope notes¶

Stellarator-symmetric states only (lasym = False), inherited from vmex.core.stability._ballooning_context().
Surfaces need iota != 0 (field-line parameterization divides by iota).
The flux tube covers one poloidal turn theta in [-pi, pi) (the solver z-grid convention of spectraxgk.core.grid.build_spectral_grid); the parallel boundary is handled by SPECTRAX-GK’s twist-shift machinery from the emitted q/s_hat/nfp.
gds21/gbdrift0 signs follow simsopt vmec_fieldlines with psi = s * psi_edge in vmex’s internal (signed) edge-flux convention. The default single-kx proxies (nx = 1) are insensitive to this overall sign, matching SPECTRAX-GK’s own VMEC bridge.

vmex.core.turbulence.GK_GEOMETRY_FIELDS = ('theta', 'gradpar', 'bmag', 'bgrad', 'gds2', 'gds21', 'gds22', 'cvdrift', 'gbdrift', 'cvdrift0', 'gbdrift0')¶: Field-line array fields of the SPECTRAX-GK flux-tube geometry contract (spectraxgk.geometry.flux_tube_contract._ARRAY_FIELDS).

vmex.core.turbulence.TURBULENCE_OBJECTIVE_NAMES = ('gamma', 'omega', 'kperp_eff2', 'linear_heat_flux_weight', 'linear_particle_flux_weight', 'mixing_length_heat_flux_proxy')¶: Ordered observables of turbulence_objective_vector() (spectraxgk.SOLVER_OBJECTIVE_NAMES).

vmex.core.turbulence.gk_fieldline_geometry(state: SpectralState, rt: SolverRuntime, *, s_index: int | None = None, alpha: float = 0.0, zeta0: float = 0.0, ntheta: int = 32, equal_arc: bool = True, arc_oversample: int = 4) → dict¶

Flux-tube geometry mapping of one field line of a converged state.

Returns the in-memory geometry contract consumed by spectraxgk.flux_tube_geometry_from_mapping (keys GK_GEOMETRY_FIELDS plus grho/jacobian and the scalar metadata q, s_hat, epsilon, R0, B0, alpha, nfp), all in the GS2/GX normalizations of simsopt vmec_fieldlines with L_ref the effective minor radius and B_ref = 2 |psi_edge| / L_ref^2 (identical to vmex.core.stability). A "vmex" sub-dict carries diagnostics used by the parity tests (dp_drho, gradpar_profile, the sampled PEST angles, …). Pure jnp — traceable and differentiable w.r.t. (state, runtime); no spectraxgk import.

Parameters:

s_index – Full-mesh surface index in [2, ns - 2]; default ~60 % of the radius (a typical core gradient region).
alpha – Field-line label alpha = theta* - iota (phi - zeta0) and the toroidal angle of the tube center.
zeta0 – Field-line label alpha = theta* - iota (phi - zeta0) and the toroidal angle of the tube center.
ntheta – Parallel samples over one poloidal turn; the emitted theta is linspace(-pi, pi, ntheta, endpoint=False) — exactly the SPECTRAX-GK solver z grid.
equal_arc – Resample the parallel coordinate so b . grad z is constant (gradpar exactly uniform, SPECTRAX-GK’s validated contract). The coordinate map is built from an arc_oversample x finer quadrature of 1/gradpar; geometry values are exact spectral evaluations at the mapped points (only the map itself is interpolated). equal_arc=False samples uniformly in the PEST angle instead (stability.py’s grid; gradpar then varies along the line and downstream use relies on spectraxgk’s mean-gradpar reduction, as in its own VMEC bridge).

vmex.core.turbulence.flux_tube_geometry(state: SpectralState, rt: SolverRuntime, *, validate: bool = False, **geometry_kwargs)¶

SPECTRAX-GK FluxTubeGeometryData for one field line (needs spectraxgk).

Thin wrapper: gk_fieldline_geometry() -> spectraxgk.flux_tube_geometry_from_mapping. validate=True turns on spectraxgk’s host-side finite/constant-gradpar checks (concrete arrays only — leave False under jit/grad tracing).

vmex.core.turbulence.turbulence_objective_vector(state: SpectralState, rt: SolverRuntime, *, selected_ky_index: int = 1, n_laguerre: int = 2, n_hermite: int = 3, nx: int = 1, ny: int = 4, lx: float = 6.0, ly: float = 12.0, params_linear=None, terms=None, r_over_lt: float | None = None, r_over_ln: float | None = None, **geometry_kwargs) → Array¶

Ordered SPECTRAX-GK linear/quasilinear observable vector (traceable).

Samples one flux tube (gk_fieldline_geometry() keyword arguments s_index/alpha/zeta0/ntheta/equal_arc pass through), builds SPECTRAX-GK’s spectral linear gyrokinetic operator on it at the selected_ky_index binormal wavenumber (ky = 2 pi k / ly in rho_ref units), selects the maximum-growth eigenbranch, and returns TURBULENCE_OBJECTIVE_NAMES (spectraxgk.solver_objective_vector_from_geometry).

The drive gradients live in SPECTRAX-GK’s LinearParams (params_linear; default: its collisionless optimization defaults R/L_n = 2.2, R/L_Ti = 6.9 — the Cyclone-base ITG drive — optionally overridden via r_over_lt/r_over_ln).

vmex.core.turbulence.turbulent_growth_rate(state: SpectralState, rt: SolverRuntime, **kwargs) → Array¶

Dominant linear gyrokinetic growth rate on one flux tube (traceable).

SPECTRAX-GK objective kind "growth": the largest real part of the eigenvalues of its spectral Hermite-Laguerre linear operator on the sampled flux tube, in v_th / L_ref units. Positive = unstable.

The operator matrix is SPECTRAX-GK’s own (spectraxgk.solver_linear_operator_matrix_from_geometry — the exact matrix behind its solver_growth_rate_from_geometry); the eigenvalue reduction here uses jnp.linalg.eigvals so the objective carries both JVP and VJP rules — SPECTRAX-GK’s dominant_real_eigenvalue is a reverse-only custom_vjp, which vmex’s forward-mode implicit Jacobian cannot trace (values agree to roundoff; gated in tests/test_turbulence.py). Keyword arguments as turbulence_objective_vector() (minus the eigenvector-dependent pieces). Two-positional (state, runtime) — drive it toward zero / negative in vmex.core.optimize.least_squares() with jac=None or jac="implicit".

vmex.core.turbulence.quasilinear_flux_proxy(state: SpectralState, rt: SolverRuntime, **kwargs) → Array¶

Mixing-length quasilinear heat-flux proxy (value-level; jac=None).

SPECTRAX-GK objective kind "quasilinear_flux": gamma * W_Q / max(kperp_eff^2, 1e-12) with W_Q the dominant mode’s normalized heat-flux weight — the mixing-length saturation rule of its quasilinear transport lane. The weight depends on the dominant eigenvector, whose non-symmetric derivatives JAX declines — use finite differences (jac=None), like the wout-engine terms. Keyword arguments as turbulence_objective_vector().

vmex.core.turbulence.nonlinear_heat_flux_proxy(state: SpectralState, rt: SolverRuntime, *, csat: float = 0.85, saturation_floor: float = 1e-10, **kwargs) → Array¶

Smooth reduced nonlinear-window heat-flux surrogate (value-level; jac=None).

SPECTRAX-GK objective kind "nonlinear_window_heat_flux": its saturation-rule closure csat * max(W_Q, 0) * 2 gamma_+ / (1 + 2.2 kperp_eff^2 + 0.15 gamma_+) mapping the linear solver row to a nonlinear heat-flux proxy (spectraxgk.objectives. vmec_transport_tables._solver_table_to_nonlinear_window_proxy — the exact objective its VMEC-JAX optimization scripts use for this kind). This is a smooth surrogate for the nonlinear transport window; SPECTRAX-GK’s docs require matched long nonlinear audits before any production nonlinear claim. Eigenvector-weighted like quasilinear_flux_proxy() — use jac=None. Keyword arguments as turbulence_objective_vector().

Outputs¶

Complete VMEC2000-compatible wout_*.nc schema, writer and reader.

This module implements the full variable set written by VMEC2000’s wrout.f (Appendix A), with the exact netCDF names, dimensions, dtypes and unit conventions of the reference implementation:

presf/pres/mass are stored in Pa (wrout.f divides the internal mu0*Pa values by mu0 on write);
jcuru/jcurv/ctor/currumnc/currvmnc are in A (again 1/mu0 applied on write);
phipf/chipf carry the twopi*signgs factor relative to the internal phips/chips arrays;
q_factor = 1/iotaf (HUGE at iota zeros);
lmns is on the half mesh (interpolated from VMEC’s internal full-mesh lambda), bsubsmns on the full mesh (converted in jxbforce.f);
lasym partner tables (rmns, zmnc, …) exist only for asymmetric runs; the free-boundary potential/surface tables (potsin, xmpot, xnpot, curlabel, *_sur) only when lfreeb.

WoutData stores every field in file convention (exactly the values found in the netCDF file), so read_wout(write_wout(x)) == x.

wout_from_state() builds the complete dataset from a converged fixed-boundary core state (vmex.core.input.VmecInput + vmex.core.solver.SpectralState), legacy-free: the Nyquist/jxbforce/Mercier tables come from the wrout.f/bss.f/ jxbforce.f/mercier.f ports in vmex.core.nyquist, the remaining VMEC2000 output quantities (eqfor.f/spectrum.f/ Compute_Currents et al.) from vmex.core.postprocess.

class vmex.core.wout.WoutData(version_: float, input_extension: str, mgrid_file: str, pcurr_type: str, pmass_type: str, piota_type: str, wb: float, wp: float, gamma: float, rmax_surf: float, rmin_surf: float, zmax_surf: float, nfp: int, ns: int, mpol: int, ntor: int, mnmax: int, mnyq: int, nnyq: int, mnmax_nyq: int, niter: int, itfsq: int, lasym: bool, lrecon: bool, lfreeb: bool, lmove_axis: bool, lrfp: bool, ier_flag: int, aspect: float, betatotal: float, betapol: float, betator: float, betaxis: float, b0: float, rbtor0: float, rbtor: float, signgs: int, IonLarmor: float, volavgB: float, ctor: float, Aminor_p: float, Rmajor_p: float, volume_p: float, ftolv: float, fsql: float, fsqr: float, fsqz: float, nextcur: int, extcur: ndarray, mgrid_mode: str, xm: ndarray, xn: ndarray, xm_nyq: ndarray, xn_nyq: ndarray, raxis_cc: ndarray, zaxis_cs: ndarray, am: ndarray, ac: ndarray, ai: ndarray, am_aux_s: ndarray, am_aux_f: ndarray, ai_aux_s: ndarray, ai_aux_f: ndarray, ac_aux_s: ndarray, ac_aux_f: ndarray, iotaf: ndarray, q_factor: ndarray, presf: ndarray, phi: ndarray, phipf: ndarray, chi: ndarray, chipf: ndarray, jcuru: ndarray, jcurv: ndarray, iotas: ndarray, mass: ndarray, pres: ndarray, beta_vol: ndarray, buco: ndarray, bvco: ndarray, vp: ndarray, specw: ndarray, phips: ndarray, over_r: ndarray, jdotb: ndarray, bdotb: ndarray, bdotgradv: ndarray, DMerc: ndarray, DShear: ndarray, DWell: ndarray, DCurr: ndarray, DGeod: ndarray, equif: ndarray, fsqt: ndarray, wdot: ndarray, rmnc: ndarray, zmns: ndarray, lmns: ndarray, gmnc: ndarray, bmnc: ndarray, bsubumnc: ndarray, bsubvmnc: ndarray, bsubsmns: ndarray, currumnc: ndarray, currvmnc: ndarray, bsupumnc: ndarray, bsupvmnc: ndarray, raxis_cs: ndarray | None = None, zaxis_cc: ndarray | None = None, rmns: ndarray | None = None, zmnc: ndarray | None = None, lmnc: ndarray | None = None, gmns: ndarray | None = None, bmns: ndarray | None = None, bsubumns: ndarray | None = None, bsubvmns: ndarray | None = None, bsubsmnc: ndarray | None = None, currumns: ndarray | None = None, currvmns: ndarray | None = None, bsupumns: ndarray | None = None, bsupvmns: ndarray | None = None, mnmaxpot: int | None = None, nobser: int | None = None, nobd: int | None = None, nbsets: int | None = None, nbfld: ndarray | None = None, potsin: ndarray | None = None, potcos: ndarray | None = None, xmpot: ndarray | None = None, xnpot: ndarray | None = None, curlabel: tuple[str, ...] | None = None, bsubumnc_sur: ndarray | None = None, bsubvmnc_sur: ndarray | None = None, bsupumnc_sur: ndarray | None = None, bsupvmnc_sur: ndarray | None = None, bsubumns_sur: ndarray | None = None, bsubvmns_sur: ndarray | None = None, bsupumns_sur: ndarray | None = None, bsupvmns_sur: ndarray | None = None)¶

Full VMEC2000 wout dataset in file conventions (wrout.f).

Every field name matches its netCDF variable name (logicals drop the __logical__ suffix). Radial arrays are (ns,); Fourier tables (ns, mnmax) / (ns, mnmax_nyq). Optional groups are None when not applicable: lasym partners for symmetric runs, free-boundary tables for fixed-boundary runs.

vmex.core.wout.write_wout(path: str | Path, data: WoutData, *, overwrite: bool = True) → Path¶

Write data to path in VMEC2000 wout_*.nc layout.

Uses NETCDF3 64-bit-offset format (as VMEC2000’s ezcdf) and reproduces wrout.f’s variable set, ordering, dimensions and attributes. extcur and mgrid_mode are created but left unwritten (netCDF fill) when nextcur == 0, matching VMEC2000.

vmex.core.wout.read_wout(path: str | Path) → WoutData¶

Read a VMEC2000-compatible wout_*.nc file into WoutData.

All values are kept in file conventions (no unit conversions), so a write_wout() / read_wout() round trip is the identity.

vmex.core.wout.wout_from_state(*, inp, state, fsqr: float, fsqz: float, fsql: float, fsqt=None, wdot=None, niter: int = 0, itfsq: int = 0, converged: bool = True, input_extension: str = '', version: float = 9.0, nextcur: int = 0, extcur=None, mgrid_mode: str = '', curlabel=None) → WoutData¶

Build a complete WoutData from a solved fixed-boundary state.

inp is the parsed vmex.core.input.VmecInput deck and state the converged vmex.core.solver.SpectralState (internal normalization, m = 1-constrained, odd-m without scalxc) — the exact representation the core solver evolves. The geometry/field state is re-evaluated on the Nyquist-extended internal grid with the core pipeline (geometry / fields), the Nyquist/jxbforce/Mercier tables come from vmex.core.nyquist, and every remaining VMEC2000 wrout.f variable from vmex.core.postprocess and the input deck.

Unlike VMEC2000 (which zeroes the late eqfor.f scalars when the run hits NITER), all derived quantities are always computed - vmex’s zero-crash policy keeps diagnostic output for non-converged states and records convergence in ier_flag (0 = converged, 2 = more iterations needed, matching vmec_params.f).

Free-boundary metadata: nextcur/extcur/mgrid_mode/ curlabel are caller-supplied (the CLI reads them from the mgrid file; extcur is the input EXTCUR array in Amperes, as wrout.f writes it). The NESTOR vacuum potential is NOT populated: potsin/ potcos/xmpot/xnpot and the *_sur surface-field tables stay None (netCDF fill) because vmex.core.freeboundary.solve_free_boundary() does not return its FreeBoundaryState (which holds potvac) — a documented gap until the solver exposes it.

vmex.core.wout.wout_field_names() → tuple[str, ...]¶: All WoutData field names (for completeness checks).

Nyquist-table, bsubs, jxbforce and Mercier post-processing (core-native).

VMEC2000 counterparts¶

Sources/Input_Output/wrout.f — the Nyquist-resolution half-mesh Fourier tables gmnc/bmnc/bsubumnc/bsubvmnc/bsupumnc/bsupvmnc/bsubsmns (plus the lasym sine partners), including the jxbforce low-pass filtering of bsubu/bsubv VMEC applies before writing them: wout_field_tables().
Sources/Input_Output/bss.f — the covariant radial field B_s on the half mesh from the metric cross terms g_su/g_sv: bsubs_half_mesh() / bsubs_full_mesh_for_wrout().
Sources/Input_Output/jxbforce.f — the low-pass filter transforms, the bsubsu/bsubsv angle derivatives and the 1D current diagnostics jdotb/bdotb/bdotgradv.
Sources/Input_Output/mercier.f (+ vmercier.f) — the Mercier stability profiles DMerc/DShear/DWell/DCurr/DGeod: mercier_and_jxb().
Sources/General/bcovar.f (IEQUI=1 block) — the surface-average correction of bsubv used for the lasym output lane: apply_bsubv_equif_correction().

The math is extracted from the parity-proven legacy modules vmex.io.wout_files.{nyquist,jxbforce,mercier,bsubs} (validated against golden VMEC2000 wout files in tests/test_wout_golden.py), re-hosted on the core types: vmex.core.geometry.RealSpaceGeometry / HalfMeshJacobian supply the parity geometry channels, vmex.core.fields.MagneticFields the half-mesh field state, and vmex.core.fourier.TrigTables (built at the Nyquist resolution) the trig/weight tables.

Normalization note (parity-critical): VMEC2000’s output transforms normalize lasym integrals on the full theta grid (dnorm = 1/(nzeta * ntheta1), fixaray.f SPH012314), while the core solver trig tables carry the reduced-interval dnorm for both symmetry modes (see vmex.core.fourier). _analysis_theta_tables() therefore rebuilds the integration-weighted theta tables with the output convention.

Host NumPy throughout: this is one-shot output post-processing, not solver code. The LBSUBS = T corrected-bsubs collocation lane of jxbforce.f (off by default in VMEC2000 and unused by every reference deck) is not ported.

class vmex.core.nyquist.WoutFieldTables(xm_nyq: ndarray, xn_nyq: ndarray, gmnc: ndarray, bmnc: ndarray, bsubumnc: ndarray, bsubvmnc: ndarray, bsupumnc: ndarray, bsupvmnc: ndarray, bsubsmns: ndarray, gmns: ndarray | None, bmns: ndarray | None, bsubumns: ndarray | None, bsubvmns: ndarray | None, bsupumns: ndarray | None, bsupvmns: ndarray | None, bsubsmnc: ndarray | None, jdotb: ndarray, bdotb: ndarray, bdotgradv: ndarray, DMerc: ndarray, DShear: ndarray, DCurr: ndarray, DWell: ndarray, DGeod: ndarray)¶

wrout.f Nyquist-resolution output tables + jxbforce/Mercier profiles.

Fourier tables are in file convention, shape (ns, mnmax_nyq); the sine partners are None for stellarator-symmetric runs. xm_nyq / xn_nyq carry the nfp factor on xn. Radial profiles are (ns,) in wout units except jdotb/bdotb (internal jxbforce.f normalization, as written by VMEC2000).

vmex.core.nyquist.nyquist_limits(trig: TrigTables) → tuple[int, int]¶

Grid Nyquist cutoffs (mnyq, nnyq) (fixaray.f).

mnyq = ntheta2 - 1 = ntheta1/2 and nnyq = nzeta/2 — geometric grid limits, independent of the retained solver modes.

vmex.core.nyquist.bsubs_half_mesh(*, geometry: RealSpaceGeometry, jacobian: HalfMeshJacobian, bsupu: ndarray, bsupv: ndarray, s: ndarray) → ndarray¶

B_s = B^u g_su + B^v g_sv on the half mesh (bss.f).

The half-mesh geometry follows jacobian.f conventions: ru12/zu12/ rs/zs come straight from half_mesh_jacobian() (identical formulas), rv12/zv12 are the half-mesh averages of the toroidal derivatives, and rs12/zs12 add the d(shalf)/ds odd-m chain-rule terms (dphids = 0.25); every half-mesh array mirrors js = 2 into the axis slot.

vmex.core.nyquist.bsubs_full_mesh_for_wrout(*, bsubs_half: ndarray) → ndarray¶: Half-mesh bsubs -> the full-mesh convention written by wrout.f.

vmex.core.nyquist.apply_bsubv_equif_correction(*, bsubv: ndarray, bsubv_e: ndarray, trig: TrigTables) → ndarray¶

bcovar.f IEQUI=1 half-mesh reconstruction of bsubv.

Backward recurrence bsubv(js) = 2*bsubv_e(js) - bsubv(js+1) followed by a per-surface constant shift preserving the surface average fpsi = <bsubv> (pwint weights, zero on the axis row).

vmex.core.nyquist.mercier_and_jxb(*, geometry: RealSpaceGeometry, s: ndarray, lasym: bool, mpol: int, ntor: int, pres: ndarray, vp: ndarray, phips: ndarray, iotas: ndarray, bsq: ndarray, sqrtg: ndarray, bsubu: ndarray, bsubv: ndarray, bsubs_half: ndarray, trig: TrigTables, signgs: int) → tuple[ndarray, ...]¶

Mercier profiles and jxbforce 1D current diagnostics.

Port of mercier.f/vmercier.f and the jdotb/bdotb/bdotgradv block of jxbforce.f, consuming the core half-mesh field state: pres in internal units (mu0*Pa), bsq the total pressure, bsubu/bsubv the jxbforce-filtered covariant fields (reduced grid for symmetric runs, full grid for lasym), bsubs_half from bsubs_half_mesh(). Returns (DMerc, DShear, DCurr, DWell, DGeod, jdotb, bdotb, bdotgradv).

vmex.core.nyquist.wout_field_tables(*, geometry: RealSpaceGeometry, jacobian: HalfMeshJacobian, metrics, fields, trig: TrigTables, s: ndarray, mpol: int, ntor: int, nfp: int, lasym: bool, vp: ndarray, phips: ndarray, iotas: ndarray, phipf: ndarray, signgs: int) → WoutFieldTables¶

Build every legacy-engine wout quantity from the core field state.

Reproduces the wrout.f output sequence: bss.f bsubs; the jxbforce.f low-pass filtering of bsubu/bsubv (with the lasym IEQUI bsubv correction feeding the antisymmetric channel); the Nyquist coefficient analysis of sqrt(g), |B|, B^u/B^v, B_u/B_v and B_s; and the jxbforce.f/mercier.f 1D diagnostics.

Inputs are the core pipeline objects evaluated on the Nyquist-extended trig tables: fields.pressure in internal units, vp/phips/iotas the half-mesh profiles (axis slots zeroed), phipf the internal full-mesh phip (as passed to magnetic_fields).

Post-processed wout quantities ported from VMEC2000 output routines.

Every function in this module is a direct numpy port of the corresponding VMEC2000 Fortran, validated against golden wout_*.nc files produced by VMEC2000 (hiddenSymmetries build, version_ = 9.0):

compute_currentsLIBSTELL/read_wout_mod.f90::Compute_Currents
(called from wrout.f); writes currumnc / currvmnc (+ *mns partners when lasym).
spectral_widthspectrum.f (with fixaray.f xmpq,
pexp = 4); writes specw.
poloidal_fluxeqfor.f (chi = twopi*chi1 accumulation);
writes chi [Wb].
safety_factor : wrout.f (qfact = 1/iotaf, HUGE at zeros).
beta_volume_profiles: eqfor.f half-mesh volume-averaged beta and
<1/R> (beta_vol, over_r, betaxis).
surface_extremaeqfor.f (rmax_surf/rmin_surf/
zmax_surf grid extrema of the boundary).
field_scalarsbcovar.f (rbtor0, rbtor from
fpsi = bvco) and eqfor.f (b0, volavgB, IonLarmor).

Unit conventions follow the netCDF file exactly (the same conventions used by wrout.f after its write-time conversions): pressures in Pa, currents in A (1/mu0 applied), phipf/chipf include the twopi*signgs factor. All inputs are file-convention arrays with shape (ns, mn) for Fourier tables and (ns,) for radial profiles.

vmex.core.postprocess.internal_angle_grid(*, ntheta: int, nzeta: int, nfp: int, lasym: bool)¶

Return VMEC’s internal angle grid and integration weights.

Mirrors fixaray.f: ntheta1 = 2*(ntheta/2), symmetric runs keep the half poloidal range [0, pi] (ntheta2 = ntheta1/2 + 1 points) with endpoint weights halved; lasym runs keep the full range (ntheta1 points) with uniform weights. zeta spans one field period.

Returns (theta, zeta, wint) with wint shaped (ntheta3, nzeta) normalized so that sum(wint) = 1.

vmex.core.postprocess.fourier_synthesis(cmn, smn, xm, xn, theta, zeta)¶

Synthesize f[js, j, k] = sum_mn cmn*cos(m u - n v) + smn*sin(...).

xm/xn are the file-convention mode tables (xn includes the nfp factor). smn may be None for stellarator-symmetric cosine series. Exactly reproduces grid values written by wrout.f because VMEC’s Nyquist coefficient tables (with the endpoint half-weight of wrout.f) are a lossless representation of the internal grid.

vmex.core.postprocess.half_mesh_r12(rmnc, rmns, xm, xn, theta, zeta)¶

Half-mesh cylindrical R exactly as bcovar.f builds r12.

VMEC stores odd-m coefficients internally divided by sqrt(s) and interpolates the even and (unscaled) odd parts separately: r12(js) = 0.5*(Re(js)+Re(js-1)) + shalf(js)*0.5*(Ro(js)+Ro(js-1)). The odd part at the magnetic axis is not recoverable from the wout coefficients (they carry a sqrt(s)=0 factor); it is linearly extrapolated, which only affects the first half-mesh surface at the ~1e-6 relative level (validated against golden VMEC2000 files).

vmex.core.postprocess.nyquist_mode_table(*, mnyq: int, nnyq: int, nfp: int)¶

VMEC2000 Nyquist mode tables (xm_nyq, xn_nyq) (fixaray.f).

Ordering: m = 0 with n = 0..nnyq, then m = 1..mnyq with n = -nnyq..nnyq; xn_nyq = n*nfp.

vmex.core.postprocess.expand_mode_columns(table, xm_old, xn_old, xm_new, xn_new)¶

Re-index a (ns, mn_old) Fourier table onto a larger mode set.

Columns present in both sets are copied; new columns are zero. Used when the solver ran with a reduced toroidal grid (ntor = 0 decks with NZETA > 1): VMEC2000 still writes the full grid-Nyquist mode set, whose extra toroidal harmonics vanish identically.

vmex.core.postprocess.compute_currents(*, bsubsmns, bsubumnc, bsubvmnc, xm_nyq, xn_nyq, bsubsmnc=None, bsubumns=None, bsubvmns=None, lasym=False)¶

Current-density harmonics currXmn = sqrt(g)*J^X (X = u, v) [A].

Port of read_wout_mod.f90::Compute_Currents (called by wrout.f). Inputs are the file-convention half-mesh bsub* tables shaped (ns, mnmax_nyq); outputs are full-mesh, already divided by mu0. Returns (currumnc, currvmnc, currumns, currvmns) with the sine partners None unless lasym.

vmex.core.postprocess.spectral_width(*, rmnc, zmns, xm, xn, rmns=None, zmnc=None, pexp=4)¶

Spectral width <M> per surface (spectrum.f; specw in wout).

specw = sum(t1*m**(pexp+1)) / sum(t1*m**pexp) with t1 the summed squares of the separable-basis (cc/ss) coefficients. In the combined helical basis of the wout tables this is rmnc**2 + zmns**2 (plus the asymmetric partners) with an extra weight 2 for n != 0 columns (the helical-to-separable change of basis doubles the sum of squares for paired +/-n modes). specw(axis) = 1.

vmex.core.postprocess.poloidal_flux(*, phips, iotas)¶

Poloidal flux chi [Wb] on the full mesh (eqfor.f).

chi = twopi * cumsum(hs * phips * iotas) with chi(axis) = 0; phips and iotas are the half-mesh wout arrays.

vmex.core.postprocess.safety_factor(iotaf)¶: q_factor = 1/iotaf with VMEC’s HUGE placeholder at zeros.

vmex.core.postprocess.mass_profile(*, pres, vp, gamma)¶

Half-mesh mass in wout units (Pa): mass = pres * vp**gamma.

VMEC evolves pres = mass / vp**gamma (bcovar.f); the wout file stores both divided by mu0. For gamma = 0 (all standard decks) mass == pres exactly, matching VMEC2000 output.

vmex.core.postprocess.beta_volume_profiles(*, bmnc, gmnc, xm_nyq, xn_nyq, pres, vp, signgs, rmnc, xm, xn, ntheta, nzeta, nfp, lasym, bmns=None, gmns=None, rmns=None)¶

Half-mesh beta_vol, betaxis and over_r (eqfor.f).

Synthesizes |B| and sqrt(g) on VMEC’s internal angular grid from the Nyquist tables (lossless), then reproduces:

tau       = signgs * wint * gsqrt
s2        = sum(bsq*tau)/vp - pres      (bsq = B^2/2 + mu0*pres)
beta_vol  = pres / s2
over_r    = sum(tau/r12) / vp
betaxis   = 1.5*beta_vol(2) - 0.5*beta_vol(3)

pres/vp are the wout half-mesh arrays (pres in Pa).

vmex.core.postprocess.surface_extrema(*, rmnc, zmns, xm, xn, ntheta, nzeta, nfp, lasym, rmns=None, zmnc=None)¶

Boundary extrema (rmax_surf, rmin_surf, zmax_surf) (eqfor.f).

VMEC evaluates the extrema on its internal grid points only (no continuous optimization), so grid synthesis reproduces the file values exactly.

vmex.core.postprocess.force_balance(*, bsubumnc, bsubvmnc, xm_nyq, xn_nyq, phipf, chipf, pres, vp, signgs)¶

Flux-surface current averages and radial force balance (fbal.f).

Recomputes, in file conventions, exactly what VMEC2000 writes:

buco/bvco: half-mesh <B_u>/<B_v> - by definition the (m,n) = (0,0) harmonic of the (parity-proven) bsub[uv]mnc tables (calc_fbal: buco = SUM(bsubu*wint));
jcuru/jcurv [A]: full-mesh surface-averaged current densities -+(signgs*ohs)*d<B_[vu]> with eqfor.f end extrapolations, /mu0;
equif: normalized radial force balance (calc_fbal + eqfor.f normalization, end points linearly extrapolated);
ctor [A]: signgs*twopi*(1.5*buco(ns) - 0.5*buco(ns1))/mu0 (bcovar.f).

phipf/chipf are the file-convention arrays (2*pi*signgs x internal); pres in Pa; vp as stored. Validated bit-exact against golden VMEC2000 wout files (symmetric and lasym).

vmex.core.postprocess.eqfor_beta_scalars(*, pres, vp, bsq, r12, bsupv, sqrtg, wint, signgs)¶

eqfor.f beta scalars (betapol, betator, betatotal).

Inputs are the internal half-mesh field state: pres in mu0*Pa, bsq the total pressure |B|^2/2 + pres, r12/bsupv/sqrtg the half-mesh grid arrays and wint the angular weights. Reproduces:

sump    = vnorm * sum(vp*pres)
sumbtot = 2*(vnorm*sum(bsq*tau) - sump)
sumbtor = vnorm*sum(tau*(r12*bsupv)**2)
beta*   = 2*sump / sumb*

vmex.core.postprocess.aspect_ratio_scalars(*, r_boundary, zu_boundary, wint)¶

aspectratio.f scalars (Aminor_p, Rmajor_p, aspect, volume_p).

r_boundary/zu_boundary are the boundary-surface R and dZ/dtheta on the internal angular grid, wint the angular weights.

vmex.core.postprocess.full_mesh_from_half(profile)¶

Full-mesh companion of a half-mesh profile (eqfor.f presf convention).

f(1) = 1.5 f(2) - 0.5 f(3), interior averages, edge extrapolation.

vmex.core.postprocess.iotaf_from_iotas(iotas)¶: add_fluxes.f90 full-mesh iota from the half-mesh profile (no RFP).

vmex.core.postprocess.chipf_from_chips(chips)¶: add_fluxes.f90 full-mesh chipf from half-mesh chips.

vmex.core.postprocess.toroidal_flux_profile(*, phipf_out, s)¶

phi [Wb] on the full mesh: rectangle-rule integral of phipf.

phi[i] = sum_{j<=i} phipf_out[j]*(s[j]-s[j-1]) with phi[0] = 0 (phipf_out in file convention, i.e. including 2*pi*signgs).

vmex.core.postprocess.lambda_wout_from_full_mesh(*, lam_full, m_modes, s, phipf_internal, lamscale)¶

Internal full-mesh lambda coefficients -> the wout half-mesh lmns.

wrout.f: the internal lambda (evolved as lamscale * lambda with the flux normalization phip) is rescaled by lamscale/phipf and interpolated to the half mesh with parity-dependent sm/sp weights; the axis row is zeroed.

vmex.core.postprocess.field_scalars(*, bvco, raxis_cc, wb, volume_p)¶

Edge/axis field scalars from bcovar.f / eqfor.f.

Returns (rbtor0, rbtor, b0, volavgB, ion_larmor):

rbtor0 = 1.5*bvco(2) - 0.5*bvco(3); rbtor same at the edge (bcovar.f with fpsi = bvco) [T*m].
b0 = rbtor0 / R_axis(v=0) with R_axis(v=0) = sum(raxis_cc) (eqfor.f b0 = fpsi0/r00) [T].
volavgB = sqrt(2*wb*(2*pi)**2 / volume_p) [T]; identical to eqfor.f sqrt(|sumbtot/volume_p|) because wb = int(B^2/2 dV)/(2*pi)**2.
IonLarmor = 0.0032/volavgB [m * sqrt(keV)].

VMEC2000-format console output.

Replicates the iteration lines, multigrid stage banners, and termination messages of VMEC2000 byte-for-column, so vmec input.x output can be diffed against xvmec2000 input.x output structurally.

VMEC2000 counterparts: Sources/Input_Output/printout.f (iteration lines, FORMATs 15/25/40/45/50/60/65/70), Sources/Initialization_Cleanup/ initialize_radial.f (FORMAT 1000, stage banner), Sources/TimeStep/ runvmec.f (FORMAT 30, force-iterations banner), and Sources/TimeStep/ eqsolve.f (FORMAT 110, vacuum banner). The exact format strings are recorded in Appendix B.

All functions return strings; the caller decides where they go (stdout, the threed1 file, or a jax.debug.callback).

vmex.core.printing.stage_banner(ns: int, mnmax: int, ftol: float, niter: int) → str¶: Multigrid stage banner (initialize_radial.f FORMAT 1000).

vmex.core.printing.vacuum_banner(iteration: int) → str¶: Free-boundary vacuum activation message (eqsolve.f FORMAT 110).

vmex.core.printing.screen_header(lasym: bool = False, lfreeb: bool = False) → str¶

Column header for the per-iteration screen line (printout.f).

Byte-exact against captured xvmec2000 output (golden fixtures): sym fixed, lasym fixed, and lasym free-boundary variants.

vmex.core.printing.screen_line(iteration: int, fsqr: float, fsqz: float, fsql: float, r_axis: float, delt: float, w_mhd: float, *, z_axis: float | None = None, del_bsq: float | None = None) → str¶

Per-iteration screen line.

printout.f FORMATs 45 (fixed sym), 50 (free sym), 65/70 (lasym): (i5,1p,3e10.2[,e11.3],e11.3,e10.2,e12.4[,e11.3]).

vmex.core.printing.threed1_header(lfreeb: bool = False) → str¶: Column header for the threed1-file iteration line (printout.f 15/25).

vmex.core.printing.threed1_line(iteration: int, fsqr: float, fsqz: float, fsql: float, fsqr_precond: float, fsqz_precond: float, fsql_precond: float, delt: float, r_axis: float, w_mhd: float, beta_vol_avg: float, spectral_width: float, *, del_bsq: float | None = None, f_edge: float | None = None) → str¶

Threed1-file iteration line (printout.f FORMAT 40).

(i6,1x,1p,7e10.2,e11.3,e12.4,e11.3,0p,f7.3,1p,2e9.2) — physical and preconditioned residuals, time step, axis position, MHD energy, volume beta, and spectral width <M>; plus vacuum diagnostics when free-boundary.

vmex.core.printing.termination_summary(ier_flag: int, input_name: str, jacobian_resets: int, total_time_s: float) → str¶

Final termination block (fileout.f).

Prints the werror message for ier_flag, the case name, the Jacobian reset count, and total wall time.

Publication-style diagnostic plots for new-core VMEC outputs (§5.1).

Self-contained matplotlib (Agg) port of the figure set from the legacy vmex.plotting module, reading everything from a wout_*.nc file (or an in-memory vmex.core.wout.WoutData):

summary 3x3 overview: iota, pressure, buco/bvco, jcuru/jcurv, DMerc, and |B| line contours at mid radius and at the plasma boundary;
surfaces flux-surface cross-sections at several zeta over one field period, with the magnetic axis marked;
modB |B| contours in (zeta, theta) at mid radius and boundary;
profiles iota / pressure / current profiles plus the fsqt force-residual convergence trace;
3d 3-D plasma boundary colored by |B|.

Both stellarator-symmetric and lasym (asymmetric) equilibria are supported: the sine/cosine partner tables (rmns, zmnc, bmns, …) are included whenever present. All figures use the Agg backend, dpi <= 150, and are closed after saving.

Public API¶

plot_wout(path_or_WoutData, outdir, which=(...)) -> dict[str, Path] plot_boozmn(path, outdir) -> dict[str, Path] plus the per-figure helpers each of those dispatches to.

vmex.core.plotting.plot_wout(wout, outdir: str | Path, which: Sequence[str] = ('summary', 'surfaces', 'modB', 'profiles', '3d'), *, name: str | None = None) → dict[str, Path]¶

Write the requested diagnostic figures for a WOUT file.

Parameters:

wout – Path to wout_*.nc or a WoutData.
outdir – Output directory (created if missing).
which – Any subset of ("summary", "surfaces", "modB", "profiles", "3d").
name – Basename prefix for the figures (default: case name from the path).
path. (Returns a mapping from figure key to the written PNG)

vmex.core.plotting.plot_boozmn(boozmn_path: str | Path, outdir: str | Path, which: Iterable[str] = ('modB', 'mode_profiles', 'spectrum'), *, name: str | None = None) → dict[str, Path]¶

Write Boozer diagnostic figures for a boozmn_*.nc file.

Returns a mapping from figure key (modB, mode_profiles, spectrum) to the written PNG path.

vmex.core.plotting.plot_summary(wout, out_path: str | Path, *, s_plot_ignore: float = 0.2) → Path¶: 3x3 overview figure (profiles + two |B| contour panels).

vmex.core.plotting.plot_surfaces(wout, out_path: str | Path, *, nzeta: int = 8, nradii: int = 8, ntheta: int = 160) → Path¶: Flux-surface cross-sections at nzeta slices over one field period.

vmex.core.plotting.plot_modB(wout, out_path: str | Path, *, ntheta: int = 90, nphi: int = 180) → Path¶: |B| contours in (phi, theta) at mid radius and the plasma boundary.

vmex.core.plotting.plot_profiles(wout, out_path: str | Path) → Path¶: Radial profiles (iota, pressure, currents) and fsqt convergence.

vmex.core.plotting.plot_boundary_3d(wout, out_path: str | Path, *, ntheta: int = 60, nzeta: int | None = None) → Path¶: 3-D plasma boundary colored by |B| (full torus).

vmex.core.plotting.plot_boozmn_modB(boozmn, out_path: str | Path, *, ntheta: int = 90, nphi: int = 180, cmap: str = 'viridis') → Path¶

Boozer-coordinate |B| contours at mid radius and the outermost surface.

cmap selects the contour colormap (pass "jet" for the STELLOPT / booz_xform convention).

vmex.core.plotting.plot_boozmn_spectrum(boozmn, out_path: str | Path, *, surface_index: int = -1, nmodes: int = 40) → Path¶: Largest Boozer |B| Fourier amplitudes on one surface (log bar chart).

vmex.core.plotting.plot_boozmn_mode_profiles(boozmn, out_path: str | Path, *, max_modes: int = 80) → Path¶: Radial Boozer |B| mode amplitudes grouped by symmetry family.

vmex.core.plotting.boozer_modB_on_surface(boozmn, *, s_index: int = -1, ntheta: int = 90, nphi: int = 180)¶

Boozer |B|(theta_B, phi_B) on one surface of a Boozer transform.

Accepts a booz_xform_jax.Booz_xform object or a boozmn_*.nc path (as produced by vmex.core.boozer.run_booz_xform()). s_index indexes the computed Boozer surfaces; -1 (the default) selects the outermost surface, i.e. |B| in Boozer coordinates on the LCFS.

Returns (theta_B, phi_B, B) where B has shape (ntheta, nphi) over one field period, suitable for a jet contour plot.

Boozer-coordinate transform driver for the new core (§9.4).

Thin, self-contained port of the booz_xform_jax invocation from the legacy vmex.booz module: read a wout_*.nc file, run the Boozer transform at the requested resolution, and write a standard boozmn_*.nc netCDF file (the writer lives inside booz_xform_jax and follows the Fortran booz_xform conventions, so downstream tools such as booz_plot and simsopt can consume the output unchanged).

Public API¶

run_booz_xform(wout_path, mbooz=32, nbooz=32, surfaces=None) -> Path: Run the transform and return the path of the written boozmn file.

Run booz_xform_jax on a WOUT file and write boozmn_*.nc.

Parameters:

wout_path – Path to a VMEC wout_*.nc file (new-core or VMEC2000 output).
mbooz – Poloidal / toroidal resolution of the Boozer spectrum.
nbooz – Poloidal / toroidal resolution of the Boozer spectrum.
surfaces – Optional surfaces to transform: values in [0, 1] select the nearest half-mesh surfaces by normalized flux, larger values are half-mesh indices. None transforms every half-mesh surface.
output_path – Explicit output file, or directory for the default boozmn_<case>.nc name (default: alongside the WOUT file).
outdir – Explicit output file, or directory for the default boozmn_<case>.nc name (default: alongside the WOUT file).
jit – Forwarded to Booz_xform.run (JIT-compiled transform kernels).
verbose – Emit booz_xform progress output.

Return type:

Path to the written boozmn netCDF file.

Return the boozmn_<case>.nc path for a WOUT file.

output_path wins outright; otherwise the file is placed in outdir (default: next to the WOUT file) with the conventional name.

Straight-axis mirrors¶

Independent paraxial fixtures for nonaxisymmetric straight mirrors.

These formulas validate equilibrium output; they do not call the mirror solver. Lengths are in metres and magnetic fields are in tesla.

class vmex.mirror.analytic.AxisymmetricPolynomialMirror(center_field: float = 1.0, half_length: float = 1.0, mirror_strength: float = 0.5)¶

Exact vacuum mirror with symmetric circular end sections.

The scalar potential is cubic in z and quadratic in transverse position. Its gradient is exactly curl-free and divergence-free, while the quartic poloidal flux gives an analytic family of nested surfaces.

property curvature: float¶: Return the quadratic axial-field coefficient in inverse metres squared.

potential(points: Any) → Any¶: Return the exact scalar magnetic potential.

field(points: Any) → Any¶: Return grad(potential) in Cartesian coordinates.

axis_field(z: Any) → Any¶: Return the quadratic on-axis mirror field.

poloidal_flux(radius: Any, z: Any) → Any¶: Return the exact axisymmetric flux labeling the nested surfaces.

boundary_radius(midplane_radius: Any, z: Any) → Any¶: Return the exact flux-surface radius at axial position z.

class vmex.mirror.analytic.FirstOrderSection(x_cos: Array, x_sin: Array, y_cos: Array, y_sin: Array)¶

First-order inverse-coordinate coefficients of an elliptical section.

x_cos: Any¶: Alias for field number 0

x_sin: Any¶: Alias for field number 1

y_cos: Any¶: Alias for field number 2

y_sin: Any¶: Alias for field number 3

class vmex.mirror.analytic.QuadrupoleField(average: Array, cosine: Array, sine: Array)¶

Coefficients in B = B0 + r^2(B20+B2c cos2a+B2s sin2a).

average: Any¶: Alias for field number 0

cosine: Any¶: Alias for field number 1

sine: Any¶: Alias for field number 2

class vmex.mirror.analytic.RotatingEllipseParaxial(half_length: float = 1.0, reference_field: float = 1.0, mirror_strength: float = 1.0, elongation: float = 2.0, rotation: float = 1.5707963267948966)¶

Flux-conserving rotating ellipse from Appendix C of Rodriguez et al.

The physical ellipse rotates by rotation between -half_length and half_length. A counter-rotation of the field-line label enforces the vacuum first-order consistency equation. The default is a 90-degree turn.

The expansion is valid for r / half_length << 1 and away from a zero of the on-axis field.

axis_field(z: Any) → Any¶: Return the even quadratic on-axis mirror field B0(z).

orientation(z: Any) → Any¶: Return the physical major-axis angle in the Cartesian cross-section.

label_angle(z: Any) → Any¶: Return the field-line-label angle required by vacuum consistency.

section_matrix(z: Any) → Any¶: Map (cos(alpha), sin(alpha)) to normalized (x, y).

first_order(z: Any) → FirstOrderSection¶: Return X1c, X1s, Y1c, Y1s at scalar axial position z.

flux_determinant(z: Any) → Any¶: Return X1c*Y1s-X1s*Y1c = Bbar/B0.

section(radius: Any, alpha: Any, z: Any) → Any¶: Return Cartesian (x,y,z) points on a first-order flux surface.

field(points: Any) → Any¶

Return the divergence-free first-order field tangent to the sections.

Transverse components follow a fixed field-line label through the section matrix. Terms beyond first order in distance from the axis are intentionally omitted, consistent with the paraxial construction.

boundary_radius(midplane_radius: Any, theta: Any, z: Any) → Any¶: Return the physical polar radius of the rotating elliptical tube.

consistency_residual(z: Any) → Any¶: Return the Appendix-C first-order vacuum consistency identity.

riccati_residual(z: Any) → Any¶: Return the Appendix-C Riccati equation residual for sigma=Y1c/Y1s.

quadrupole(z: Any) → QuadrupoleField¶: Return the second-order field coefficients from Appendix-C Eq. (80).

center_quadrupole() → QuadrupoleField¶: Return the independent closed form at the magnetic-well minimum.

field_strength(radius: Any, alpha: Any, z: Any) → Any¶: Evaluate the paraxial field strength through order r^2.

class vmex.mirror.analytic.StraightFieldLineMirror(center_field: float = 1.0, axial_scale: float = 2.0)¶

Agren-Savenko marginal-minimum-B vacuum field through paraxial order.

The scalar potential is their Eq. (2), truncated at relative order (radius/axial_scale)^4. Use only for |z| < axial_scale and a thin flux tube.

potential(points: Any) → Any¶: Return the paraxial scalar magnetic potential at Cartesian points.

field(points: Any) → Any¶: Return grad(potential) in Cartesian coordinates.

axis_field(z: Any) → Any¶: Return the exact on-axis field B0/(1-z^2/c^2).

clebsch_labels(points: Any) → Any¶: Return leading-order straight-line labels (x0,y0).

field_line(x0: Any, y0: Any, z: Any) → Any¶: Return leading-order straight, nonparallel field-line points.

section(midplane_radius: Any, alpha: Any, z: Any) → Any¶: Return an analytic flux-tube section from a circular midplane.

boundary_radius(midplane_radius: Any, theta: Any, z: Any) → Any¶: Return the physical polar radius of its elliptical flux tube.

ellipticity(z: Any) → Any¶: Return the major/minor axis ratio of the analytic section.

Cubic B-spline coefficients for open and closed mirror equilibria.

Knot locations are static NumPy data; coefficient evaluation and transfer are differentiable JAX operations.

class vmex.mirror.splines.ClosedFieldLine(axial_parameter: Any, theta: Any, iota: Any)¶: One field line traced through a periodic spline equilibrium.

class vmex.mirror.splines.CubicBSplineBasis(knots: ndarray, breakpoints: ndarray, periodic: bool, size: int, collocation_nodes: ndarray, quadrature_nodes: ndarray, quadrature_weights: ndarray)¶

Static clamped or periodic cubic basis with JAX evaluation.

classmethod clamped(breakpoints: Any, *, quadrature_order: int = 4) → CubicBSplineBasis¶: Build an open cubic basis with repeated endpoint knots.

classmethod periodic_uniform(size: int, domain: tuple[float, float] = (0.0, 6.283185307179586), *, quadrature_order: int = 4) → CubicBSplineBasis¶: Build folded uniform cubic splines on a periodic interval.

property domain: tuple[float, float]¶: Return the open or fundamental periodic interval.

basis_matrix(points: Any, *, derivative: int = 0) → Any¶: Evaluate basis values or derivatives at one-dimensional points.

evaluate(coefficients: Any, points: Any, *, derivative: int = 0, axis: int = -1) → Any¶: Evaluate spline coefficients along axis at arbitrary points.

fit(values: Any, *, nodes: Any | None = None, axis: int = -1) → Any¶: Interpolate nodal values to coefficients with a square collocation solve.

integrate(coefficients: Any, *, axis: int = -1) → Any¶: Integrate a spline using per-span Gauss-Legendre quadrature.

insert_knot(coefficients: Any, knot: float, *, axis: int = -1) → tuple[CubicBSplineBasis, Any]¶: Insert one open knot exactly with the Boehm coefficient update.

refine_periodic_uniform(coefficients: Any, target_size: int, *, axis: int = -1) → tuple[CubicBSplineBasis, Any]¶: Exactly refine a uniform periodic cubic spline by dyadic subdivision.

class vmex.mirror.splines.QIMirrorSplice(points: ndarray, leg_windows: tuple[tuple[float, float], ...], total_length: float, corner_angle: float, closure_error: float, leg_directions: ndarray, leg_lengths: ndarray, cut_points: ndarray, cut_indices: tuple[int, ...])¶

Closed magnetic axis formed by inserting straight mirror legs.

points is a closed, arc-length-orderable polyline (the final point precedes the first). A quasi-isodynamic axis is cut at its low-curvature symmetry planes (2 * nfp of them – four for nfp = 2) and at each cut an exactly-straight mirror leg is inserted along the local axis tangent, so every leg continues the axis in its own direction. The per-cut leg lengths (leg_lengths, one per cut) are chosen so the inserted displacements cancel (the loop closes), and the curve is assembled from one stellarator-symmetric half reflected 180 degrees about the x axis, so the result is stellarator symmetric to rounding.

leg_windows are the arc-length spans of the straight legs (one per cut, in curve order). corner_angle (degrees) is the residual tangent break at the leg/return junctions – now near zero, since each leg is tangent to the axis: the seam is a curvature break (the exactly-zero-curvature leg meeting the finite-curvature return), which a global Fourier basis still rings on and a local B-spline represents cleanly. closure_error is the residual of the closed loop (zero to rounding by construction).

class vmex.mirror.splines.SplineMirrorBoundary(radius_coefficients: Any)¶: Lateral mirror boundary stored as axial B-spline coefficients.

class vmex.mirror.splines.SplineMirrorDiscretization(spline: CubicBSplineBasis, grid: MirrorGrid, evaluation_matrix: ndarray, closed: bool = False)¶

Coefficient-to-quadrature map for a fixed mirror configuration.

classmethod build(config: MirrorConfig, *, elements: int, quadrature_order: int = 4) → SplineMirrorDiscretization¶: Build a clamped spline and endpoint-augmented Gauss mirror grid.

classmethod build_cgl(config: MirrorConfig, *, elements: int, quadrature_order: int = 4) → SplineMirrorDiscretization¶: Build coefficients evaluated on the CGL grid used by exterior panels.

classmethod build_closed(resolution: MirrorResolution, *, coefficient_count: int, quadrature_order: int = 4) → SplineMirrorDiscretization¶: Build a periodic longitudinal spline grid for a closed hybrid.

evaluate_boundary(boundary: SplineMirrorBoundary) → MirrorBoundary¶: Evaluate boundary coefficients on the solver quadrature grid.

evaluate_state(state: SplineMirrorState) → MirrorState¶: Evaluate state coefficients on the solver quadrature grid.

fit_boundary(boundary: MirrorBoundary, source_grid: MirrorGrid) → SplineMirrorBoundary¶: Fit a nodal boundary once to initialize coefficient-native solves.

fit_state(state: MirrorState, source_grid: MirrorGrid) → SplineMirrorState¶: Fit a nodal state once to initialize coefficient-native solves.

project_fixed_boundary(state: SplineMirrorState, boundary: SplineMirrorBoundary) → SplineMirrorState¶

Apply side geometry, axis regularity, and the lambda gauge.

Open endpoint coefficients remain prescribed cut profiles. Periodic states have no distinguished end coefficients.

impose_self_similar_cuts(state: SplineMirrorState, boundary: SplineMirrorBoundary) → SplineMirrorState¶

Fix both end cuts to scaled copies of their LCFS sections.

radius_scale is independent of s on a self-similar cut, while the physical radius still scales as sqrt(s). The stream function remains supplied by the field initializer.

transfer_boundary(state: SplineMirrorState, source: SplineMirrorBoundary, target: SplineMirrorBoundary) → SplineMirrorState¶: Rescale a nested restart from source to target boundary.

class vmex.mirror.splines.SplineMirrorSolveResult(coefficient_state: SplineMirrorState, evaluated: Any)¶: Converged coefficient state and its evaluated mirror result.

class vmex.mirror.splines.SplineMirrorState(radius_coefficients: Any, lambda_coefficients: Any)¶: Geometry and stream-function B-spline coefficients.

class vmex.mirror.splines.StellaratorMirrorSetup(discretization: Any, axis_coefficients: Any, axis: Any, boundary: SplineMirrorBoundary, initial_state: SplineMirrorState)¶: Periodic spline discretization, axis, LCFS, and nested initial state.

vmex.mirror.splines.build_qi_mirror_hybrid(axis_points: Any, resolution: MirrorResolution, *, cut_indices: tuple[int, ...], straight_length: float, section_radius: float, coefficient_count: int = 64, axial_flux_derivative: Any = 0.02, quadrature_order: int = 3, samples_per_leg: int = 256) → StellaratorMirrorSetup¶

Build a solvable closed-spline hybrid from a spliced QI magnetic axis.

The closed axis_points are spliced with splice_straight_legs(), fitted to a periodic cubic B-spline of coefficient_count controls at uniform arc length, and wrapped with a constant circular cross-section of radius section_radius. A circular section is rotation-invariant, so the large frame holonomy of a fully three-dimensional QI axis does not enter the boundary. The returned StellaratorMirrorSetup feeds solve_fixed_boundary() exactly like the analytic racetrack.

vmex.mirror.splines.build_stellarator_mirror_hybrid(resolution: MirrorResolution, *, coefficient_count: int = 32, axis_coefficient_count: int | None = None, straight_length: float = 8.0, return_radius: float = 2.5, semi_major: float = 0.45, semi_minor: float = 0.3, section_turns: int = 0, axial_flux_derivative: Any = 0.02, quadrature_order: int = 4) → StellaratorMirrorSetup¶

Build a closed two-leg mirror with rotating stellarator returns.

section_turns turns the elliptical cross-section continuously around the closed circuit by that many full 2*pi turns, superposed on the return-only 90-degree rotation. The straight legs keep an exactly straight axis; only the ellipse they carry rotates, so a nonzero section_turns raises the rotational transform (by amplifying the current-driven transform: iota 0.085 -> 0.141 at s=0.75 for two turns). The default 0 reproduces the legacy return-only rotation.

axis_coefficient_count freezes the leg-return junction transition. The junction between an exactly straight leg (zero curvature) and a circular return (curvature 1/return_radius) is rounded across the cubic spline’s local support, so building the axis directly in a finer solve basis narrows that transition and sharpens the curvature overshoot at the junction as fast as refinement helps. When axis_coefficient_count is set, the racetrack axis and rotating section are constructed at that base control count and then exactly refined (refine_periodic_uniform) to coefficient_count, so the junction-transition width is a fixed design parameter of the family while the equilibrium solve basis refines. coefficient_count must be a dyadic multiple of axis_coefficient_count. The default None keeps the legacy behaviour of building the geometry directly in the solve basis.

vmex.mirror.splines.solve_fixed_boundary(initial_state: SplineMirrorState, boundary: SplineMirrorBoundary, discretization: SplineMirrorDiscretization, config: MirrorConfig, *, axial_flux_derivative: Any, mass_profile: Any = 0.0, current_derivative: Any = 0.0, gamma: float = 1.6666666666666667, solve_lambda: bool = False, axis: Any | None = None, gradient_tolerance: float = 1e-11, require_convergence: bool = False) → SplineMirrorSolveResult¶: Solve an open or periodic closed scalar-pressure equilibrium.

vmex.mirror.splines.solve_fixed_boundary_from_radius(radius: Any, config: MirrorConfig, *, elements: int = 6, axial_flux_derivative: Any = 0.01, solve_lambda: bool = True, require_convergence: bool = True, **kwargs: Any) → SplineMirrorSolveResult¶

Solve a fixed-boundary mirror equilibrium directly from a boundary radius.

This is the one-call convenience entry point. It builds the source grid, boundary, spline discretization, fitted boundary, and self-similar initial state, then runs solve_fixed_boundary(). radius may be a scalar, an axial (nxi,) array, or a full (ntheta, nxi) array (see MirrorBoundary.from_radius()). Extra keyword arguments are forwarded to solve_fixed_boundary(); power users who need explicit control over the discretization or initial state should call that function directly.

vmex.mirror.splines.splice_straight_legs(axis_points: Any, *, cut_indices: tuple[int, ...], straight_length: float, samples_per_leg: int = 256) → QIMirrorSplice¶

Cut a stellarator-symmetric closed axis and insert tangent-aligned legs.

The closed input axis_points (shape (P, 3)) is cut at the N cut_indices (its low-curvature symmetry planes – 2 * nfp of them), and at each cut an exactly-straight mirror leg is inserted along the local axis tangent, so every leg continues the axis in its own direction (no shared “bisector”, so each leg meets its curved returns tangentially). The per-cut leg lengths are set by _closure_leg_lengths() so the inserted displacements cancel, and the curve is built from one stellarator-symmetric half (between the two symmetry-fixed cuts) reflected 180 degrees about the x axis, so the racetrack is stellarator symmetric to rounding.

Requires a stellarator-symmetric axis whose cuts include exactly two symmetry-fixed planes (where the tangent reverses under the symmetry); an nfp = 2 QI axis cut at its four low-curvature planes satisfies this.

vmex.mirror.splines.trace_closed_field_line(field: Any, discretization: SplineMirrorDiscretization, *, radial_index: int, theta0: float = 0.0, turns: int = 1, steps_per_turn: int = 256) → ClosedFieldLine¶: Integrate dtheta/du = B^theta/B^u with periodic RK4 steps.

Mirror input contracts and differentiable state containers.

The supported open-end model is a finite equilibrium domain between two fixed, flux-carrying cuts. These cuts are not periodic and are not plasma-vacuum interfaces. The lateral s=1 surface is the fixed or free plasma boundary.

class vmex.mirror.model.MirrorResolution(ns: int = 17, mpol: int = 0, nxi: int = 33)¶

Static resolution for (s, theta, xi) mirror coordinates.

mpol is the largest retained theta Fourier mode. Axisymmetry uses mpol=0. The collocation size is derived as 2*mpol+1 so no undeclared or Nyquist mode enters the state.

property ntheta: int¶: Number of nodal values required to represent modes through mpol.

property axisymmetric: bool¶: Whether theta dependence is absent.

class vmex.mirror.model.MirrorConfig(resolution: MirrorResolution = MirrorResolution(ns=17, mpol=0, nxi=33), z_min: float = -1.0, z_max: float = 1.0, ftol: float = 1e-12, max_iterations: int = 2000)¶

Numerical and boundary contract for a mirror equilibrium.

Geometry and normal flux are fixed at both axial cuts while field lines may cross them. End losses, sheaths, sources, and transport are outside this equilibrium model.

The default nonlinear tolerance is the requested component-wise physical force tolerance. It is not an optimizer objective tolerance.

build_grid() → MirrorGrid¶: Build immutable collocation and quadrature data.

class vmex.mirror.model.MirrorBoundary(radius_scale: Any)¶

Lateral boundary scale a(theta, xi) in r=sqrt(s)*a.

radius_scale has shape (ntheta, nxi). It is a differentiable JAX leaf so fixed-boundary shape derivatives do not require another boundary representation.

classmethod from_radius(radius: Array, grid: MirrorGrid) → MirrorBoundary¶: Broadcast scalar, axial, or full theta-axial radii to the grid.

classmethod from_axis_field(axial_flux_derivative: Array, on_axis_bz: Array, grid: MirrorGrid, *, radius_floor: float = 0.0) → MirrorBoundary¶

Build the leading-order flux tube a=sqrt(2*Psi'/|Bz|).

This paraxial relation is an initializer and analytic validation fixture, not a replacement for a finite-radius equilibrium solve.

class vmex.mirror.model.MirrorState(radius_scale: Any, lambda_stream: Any)¶

Differentiable mirror geometry and field-line state.

The first two arrays have shape (ns, ntheta, nxi). radius_scale defines r=sqrt(s)*radius_scale; storing the regular scale rather than r avoids evolving a singular radial derivative at the magnetic axis. lambda_stream is the divergence-free field stream function and uses a zero surface-average gauge in the solver lane.

classmethod from_boundary(boundary: MirrorBoundary, grid: MirrorGrid) → MirrorState¶: Construct the radial self-similar initial state for a boundary.

validate_shape(grid: MirrorGrid) → None¶: Raise when state arrays do not match the static grid.

vmex.mirror.model.project_fixed_boundary_state(state: MirrorState, boundary: MirrorBoundary, grid: MirrorGrid) → MirrorState¶

Apply the side boundary, axis regularity, and lambda gauge.

The input state’s endpoint profiles are prescribed cut data. Keeping them intact permits finite-radius flux surfaces instead of forcing every cut to be a scaled copy of the LCFS. The lambda surface mean is a pure gauge.

Shared host optimization and preconditioning for mirror equilibria.

The coefficient fixed-boundary solver and the nodal free-boundary solver use the same convergence contract: optimizer status never substitutes for a normalized variational residual below MirrorConfig.ftol. Open systems use matrix-free Newton-GMRES at every size, with a bounded dense rescue for small stalled systems.

class vmex.mirror.solver.MirrorSolveResult(state: MirrorState, energy: MirrorEnergy, variational: VariationalResidual, force: IsotropicForceResidual, staggered_weak_force: VariationalResidual | None, normalized_divergence_rms: Any, history: Any, iterations: int, converged: bool, optimizer_success: bool, linear_iterations: int, final_linear_residual: float, message: str)¶

Solved state, diagnostics, and dense iteration history.

History columns are iteration, total_energy, radius_variational_rms, lambda_variational_rms, variational_max. The independently reconstructed pointwise force is stored in force and evaluated for the final state. optimizer_success is recorded separately and never substitutes for converged.

exception vmex.mirror.solver.MirrorConvergenceError(result: MirrorSolveResult)¶

Raised when a caller requires a physically converged mirror solve.

Subclasses vmex.core.errors.VmecError so mirror solve failures flow through the same zero-crash typed-error taxonomy as the core solver.

class vmex.mirror.solver.SeparableMirrorPreconditioner(radial_vectors: ndarray, axial_vectors: ndarray, denominator: ndarray, active_shape: tuple[int, int, int])¶

Tensor-product inverse for radial, poloidal, and axial stiffness.

The model is a shifted radial/Fourier/CGL stiffness sum. Its inverse uses two small eigendecompositions and FFTs, not a dense 3D factorization.

classmethod build(grid: MirrorGrid, *, radial_nodes: int | None = None, shift: float = 0.001, radial_strength: float = 1.0, poloidal_strength: float = 1.0, axial_strength: float = 1.0) → SeparableMirrorPreconditioner¶: Build the normalized separable stiffness inverse for grid.

classmethod build_from_axial_stiffness(grid: MirrorGrid, axial_stiffness: Array, *, radial_nodes: int | None = None, poloidal_nodes: int | None = None, shift: float = 0.001, radial_strength: float = 1.0, poloidal_strength: float = 1.0, axial_strength: float = 1.0) → SeparableMirrorPreconditioner¶: Build the tensor inverse from a representation-specific axial block.

property size: int¶: Number of active geometry unknowns.

apply(vector: Any) → Any¶: Apply the inverse model operator to a flattened residual.

operator(vector: Any) → Any¶: Apply the model operator, primarily for verification tests.

apply_gauge_free(vector: Any, *, free_indices: ndarray, pivot: int, weights: ndarray) → Any¶

Apply the inverse after eliminating one weighted-mean gauge node.

The lift and orthogonal projection keep the reduced operator symmetric.

Coupled plasma-boundary-vacuum solves for straight-axis mirrors.

class vmex.mirror.free_boundary.FreeBoundaryParameters(external_field: Any, axial_flux_derivative: Any, mass_profile: Any, current_derivative: Any)¶: Differentiable physical controls for a free-boundary equilibrium.

class vmex.mirror.free_boundary.FreeBoundaryMirrorResult(coefficient_boundary: SplineMirrorBoundary, coefficient_state: SplineMirrorState, boundary: MirrorBoundary, plasma_state: MirrorState, plasma_energy: MirrorEnergy, plasma_force: IsotropicForceResidual, plasma_staggered_weak_force: VariationalResidual | None, normalized_divergence_rms: Array, plasma_b_squared: Array, pressure: Array, vacuum_geometry: ClosedMirrorSurface, vacuum_field: ExteriorVacuum, mass_scale: Array, plasma_scale: float, target_central_pressure: float | None, interface: InterfaceResidual, history: Array, variational_max: Array, iterations: int, linear_iterations: int, final_linear_residual: float, converged: bool, optimizer_success: bool, message: str)¶: Joint plasma-boundary-vacuum equilibrium result.

vmex.mirror.free_boundary.solve_free_boundary(initial_boundary: SplineMirrorBoundary, discretization: SplineMirrorDiscretization, config: MirrorConfig, external_field: Any, *, axial_flux_derivative: Any, mass_profile: Any = 0.0, current_derivative: Any = 0.0, solve_lambda: bool | None = None, gamma: float = 1.6666666666666667, initial_state: SplineMirrorState | None = None, target_central_pressure: float | None = None, initial_mass_scale: float = 1.0, exterior_ntheta: int = 40, exterior_order: int = 8, exterior_spectral_side_density: bool = False, require_convergence: bool = False) → FreeBoundaryMirrorResult¶: Solve an open free-boundary mirror in longitudinal B-spline coefficients.

vmex.mirror.free_boundary.solve_beta_scan(initial_boundary: SplineMirrorBoundary, discretization: SplineMirrorDiscretization, config: MirrorConfig, external_field: Any, beta_values: Any, *, axial_flux_derivative: Any, reference_field: float, current_derivative: Any = 0.0, gamma: float = 1.6666666666666667, beta_rtol: float = 1e-08, initial_state: SplineMirrorState | None = None, initial_restart: FreeBoundaryRestart | None = None, exterior_ntheta: int = 40, exterior_order: int = 8, exterior_spectral_side_density: bool = False) → tuple[FreeBoundaryMirrorResult, ...]¶: Continue one coefficient-native free-boundary state through increasing beta.

Implicit coefficient-fixed and free-boundary mirror derivatives.

Derivatives use converged equilibrium equations, not host iteration histories.

class vmex.mirror.implicit.FreeBoundaryAdjointConfig(axisymmetric_ntheta: int = 40, exterior_order: int = 8, spectral_side_density: bool = False, gamma: float = 1.6666666666666667, rtol: float = 1e-09, max_restarts: int = 30)¶: Static exterior quadrature and linear-solver controls.

class vmex.mirror.implicit.FreeBoundaryAdjointResult(value: Any, gradient: FreeBoundaryParameters, iterations: int, relative_residual: float, converged: bool)¶: Scalar value, total control gradient, and transpose-solve diagnostics.

class vmex.mirror.implicit.FreeBoundaryParameters(external_field: Any, axial_flux_derivative: Any, mass_profile: Any, current_derivative: Any)¶: Differentiable physical controls for a free-boundary equilibrium.

class vmex.mirror.implicit.MirrorAdjointResult(value: Any, gradient: Any, iterations: int, relative_residual: float, converged: bool, linear_solver: str)¶: Scalar value, total parameter gradient, and linear-solve diagnostics.

class vmex.mirror.implicit.MirrorTangentResult(tangent: MirrorState, iterations: int, relative_residual: float, converged: bool)¶: Equilibrium state tangent and linear-solve diagnostics.

class vmex.mirror.implicit.SplineFixedBoundaryParameters(boundary_coefficients: Any, axis_coefficients: Any, axial_flux_derivative: Any, mass_profile: Any, current_derivative: Any)¶: Differentiable controls for a coefficient-native spline equilibrium.

vmex.mirror.implicit.free_boundary_adjoint(result: Any, parameters: FreeBoundaryParameters, discretization: Any, quantity: Callable[[MirrorBoundary, MirrorState, MirrorEnergy, Any], Any], *, config: FreeBoundaryAdjointConfig = FreeBoundaryAdjointConfig(axisymmetric_ntheta=40, exterior_order=8, spectral_side_density=False, gamma=1.6666666666666667, rtol=1e-09, max_restarts=30)) → FreeBoundaryAdjointResult¶

Differentiate a scalar through a converged axisymmetric exterior solve.

End-cut radii are fixed. Pressure-profile, flux, current, applied-field, and solved lateral-shape responses follow the same primal residual.

vmex.mirror.implicit.spline_fixed_boundary_adjoint(result: Any, parameters: SplineFixedBoundaryParameters, discretization: Any, quantity: Callable[[MirrorState, MirrorEnergy], Any], *, gamma: float = 1.6666666666666667, solve_lambda: bool = False, rtol: float = 1e-10, max_restarts: int = 20) → MirrorAdjointResult¶: Differentiate a scalar through a converged spline equilibrium.

vmex.mirror.implicit.spline_fixed_boundary_tangent(result: Any, parameters: SplineFixedBoundaryParameters, parameter_tangent: SplineFixedBoundaryParameters, discretization: Any, *, gamma: float = 1.6666666666666667, solve_lambda: bool = False, rtol: float = 1e-10, max_restarts: int = 20) → MirrorTangentResult¶: Differentiate a converged spline state in one control direction.

vmex.mirror.implicit.spline_fixed_boundary_parameters(boundary: Any, *, axial_flux_derivative: Any, mass_profile: Any = 0.0, current_derivative: Any = 0.0, axis_coefficients: Any | None = None) → SplineFixedBoundaryParameters¶: Collect native boundary, axis, flux, pressure, and current controls.

Mirror-native outputs and compact restart files.

mout is deliberately separate from VMEC’s toroidal wout schema. It stores physical-grid arrays so a solved open-ended equilibrium can be plotted or inspected without reconstructing the solver objects.

class vmex.mirror.output.AxisymmetricBetaDiagnostics(requested_beta: Any, achieved_reference_beta: Any, volume_averaged_beta: Any, local_axis_beta: Any, center_radius: Any, center_axis_field: Any, center_vacuum_side_field: Any, diamagnetic_field_ratio: Any, paraxial_field_ratio: Any, paraxial_relative_error: Any)¶: Scalar checks for one axisymmetric free-boundary beta point.

class vmex.mirror.output.FreeBoundaryRestart(boundary: SplineMirrorBoundary, plasma_state: SplineMirrorState, mass_scale: float)¶

Coefficient-native state needed to hot-start a free-boundary solve.

classmethod from_result(result: Any) → FreeBoundaryRestart¶: Extract restart data from a free-boundary solve result.

class vmex.mirror.output.MoutData(s: Any, theta: Any, xi: Any, z: Any, boundary_radius: Any, radius_scale: Any, lambda_stream: Any, mod_b: Any, b_xyz: Any, pressure: Any, history: Any, coil_xyz: Any, ftol: float, iterations: int, converged: bool, mass_scale: float, variational_max: float, normal_stress_rms: float, b_normal_rms: float, staggered_weak_max: float = nan, pointwise_force_rms: float = nan, normalized_divergence_rms: float = nan, message: str = '', schema: str = 'vmex.mirror.mout/1')¶: Data-only representation of one straight-axis mirror equilibrium.

vmex.mirror.output.load_free_boundary_restart(path: str | Path, discretization: SplineMirrorDiscretization) → FreeBoundaryRestart¶: Load coefficient-native restart data for discretization.

vmex.mirror.output.mout_from_result(result: Any, grid: Any, config: Any, *, axial_flux_derivative: Any, current_derivative: Any = 0.0, boundary: Any | None = None, pressure: Any | None = None, coil_xyz: Any | None = None) → MoutData¶: Collect a fixed- or free-boundary result and plotting fields.

vmex.mirror.output.plot_axisymmetric_beta_scan_summary(entries, outdir: str | Path, *, display: tuple[int, ...], name: str = 'mirror_free_boundary_beta50_summary', strong_force_gate: float | None = None) → Path¶

Render one tight beta-scan composite: 3D states plus scan diagnostics.

entries is a sequence of (label, mout, supported) tuples in increasing-beta order; display selects which entries get a 3D panel. All entries appear in the shared diagnostics row. Unsupported (validation-only) entries are drawn dashed.

Render two solved fixed-boundary mirrors side by side in 3D.

Each panel is coloured by its own LCFS |B| range with an attached colorbar, so an axisymmetric mirror and a rotating-ellipse mirror can be compared at a glance.

vmex.mirror.output.plot_mout(mout: MoutData | str | Path, outdir: str | Path, *, name: str | None = None) → dict[str, Path]¶: Render summary, cross-section, |B|, and horizontal 3D plots.

vmex.mirror.output.plot_stellarator_mirror_hybrid(result: Any, setup: Any, outdir: str | Path, *, name: str = 'stellarator_mirror_hybrid') → Path¶: Plot a solved periodic two-mirror/stellarator hybrid equilibrium.

vmex.mirror.output.read_mout(path: str | Path) → MoutData¶: Read a MoutData file and validate its schema.

vmex.mirror.output.save_free_boundary_restart(path: str | Path, restart: FreeBoundaryRestart | Any) → Path¶: Atomically write a compressed, data-only free-boundary restart.

vmex.mirror.output.summarize_axisymmetric_beta_scan(results: tuple['FreeBoundaryMirrorResult', ...], requested_betas: Any, grid: MirrorGrid, *, reference_field: float) → tuple[AxisymmetricBetaDiagnostics, ...]¶

Summarize solved beta points against the beta-zero equilibrium.

achieved_reference_beta uses the supplied vacuum reference field, while local_axis_beta uses the finite-beta plasma field. The paraxial comparison is B/B_vac = sqrt(1-beta) and is meaningful for a long, approximately cylindrical mirror away from beta one.

vmex.mirror.output.write_mout(path: str | Path, data: MoutData, *, overwrite: bool = True) → Path¶: Write a compact mirror-native NetCDF file.

Errors and CLI¶

Typed exception taxonomy for vmex (zero-crash policy).

Every physics or input failure maps to one of these exceptions instead of a crash, a bare traceback, or sys.exit. Each exception carries the diagnostic state needed to understand the failure (iteration counters, force residuals, offending surface). The CLI catches VmecError and prints the VMEC2000-style termination message from WERROR_MESSAGES plus a one-line remedy hint.

VMEC2000 counterpart: the ier_flag error codes defined in Sources/General/vmec_params.f and the werror message table printed by Sources/Input_Output/fileout.f. See §2.5.

vmex.core.errors.WERROR_MESSAGES: dict[int, str] = {0: 'EXECUTION TERMINATED NORMALLY', 1: 'INITIAL JACOBIAN CHANGED SIGN!', 2: 'MORE ITERATIONS REQUIRED', 4: 'MORE THAN 75 JACOBIAN ITERATIONS (DECREASE DELT)', 5: 'ERROR READING INPUT FILE OR NAMELIST', 7: 'PHIEDGE HAS WRONG SIGN IN VACUUM REGION', 8: 'NS ARRAY MUST NOT BE ALL ZEROES', 9: 'ERROR IN INPUT VALUES', 11: 'EXECUTION TERMINATED NORMALLY'}¶: VMEC2000 termination messages, keyed by ier_flag (Sources/Input_Output/fileout.f, werror table).

exception vmex.core.errors.VmecError(message: str, hint: str = '', ier_flag: int = 9)¶

Base class for all vmex failures.

message¶

Human-readable description (VMEC2000-style where applicable).

Type:: str

hint¶

One-line remedy suggestion shown by the CLI.

Type:: str

ier_flag¶

The matching VMEC2000 ier_flag code, for wout/status parity.

Type:: int

exception vmex.core.errors.VmecInputError(message: str, hint: str = '', ier_flag: int = 5)¶

Invalid or unreadable input (INDATA / JSON / arguments).

VMEC2000: input_error_flag paths in readin.f.

exception vmex.core.errors.VmecJacobianError(message: str, hint: str = '', ier_flag: int = 4, iteration: int = 0, jacobian_resets: int = 0, fsq: tuple[float, float, float] | None = None)¶

The flux-surface Jacobian changed sign and could not be recovered.

Raised after the VMEC2000 escalation ladder is exhausted (axis re-guess, time-step resets at ijacob = 25/50, abort at 75 → jac75_flag). VMEC2000: Sources/General/jacobian.f (irst=2) and eqsolve.f.

exception vmex.core.errors.VmecConvergenceError(message: str, hint: str = '', ier_flag: int = 2, iteration: int = 0, fsq: tuple[float, float, float] | None = None, ftol: float = 0.0)¶

The force residuals did not reach ftol within the iteration budget.

VMEC2000: more_iter_flag from eqsolve.f. Carries the residual history tail so callers can decide whether to continue (hot restart).

exception vmex.core.errors.MgridNotFoundError(message: str, hint: str = '', ier_flag: int = 5, path: str = '')¶

A free-boundary run referenced an mgrid file that cannot be read.

The solver catches this and falls back to a fixed-boundary solve with a warning (behavior VMEC2000 has and VMEC++ dropped — §2.5); it is re-raised only when the caller explicitly requires free-boundary.

vmex command-line entry point (new core, fixed-boundary).

VMEC2000 counterpart: the xvmec2000 executable driver (Sources/TimeStep/vmec.f / runvmec.f): parse the input deck, run the NS_ARRAY multigrid ladder with VMEC2000-format console output, write the wout_<case>.nc file, and print the fileout.f termination summary.

The solve path is the clean-room core end to end: vmex.core.input.VmecInput (INDATA or VMEC++-style JSON) -> vmex.core.multigrid.solve_multigrid() -> vmex.core.wout.wout_from_state() -> vmex.core.wout.write_wout(), plus the core plotting (--plot) and Boozer (--booz) drivers.

Zero-crash policy (§2.5): every failure maps to a typed vmex.core.errors.VmecError; the CLI prints the VMEC2000 werror message plus a one-line hint and exits with the matching ier_flag code.

Free-boundary routing (LFREEB = T):

a readable mgrid file goes to vmex.core.freeboundary.solve_free_boundary() with the VMEC2000 console output (In VACUUM block, VACUUM PRESSURE TURNED ON banner) and free-boundary wout metadata (nextcur/extcur/curlabel/ mgrid_mode);
a missing mgrid file falls back to a fixed-boundary solve with a warning (VMEC2000 behavior, dropped by VMEC++);
MGRID_FILE = 'DIRECT_COILS' (or the --coils flag) builds the external field from an ESSOS coils file (essos.coils.Coils): the coils are tabulated into an in-memory mgrid (Coils.to_mgrid) and read back as an vmex.core.mgrid.MgridField (solve_free_boundary(inp, external_field=mgrid_field)); requires ESSOS.

Documented divergences of the free-boundary lane:

solve_free_boundary() is single-grid with no warm-start seam, so only the final NS_ARRAY stage is run (from the standard interior guess); VMEC2000 runs the whole ladder with vacuum re-activating per stage. Multi-stage decks print a note.
The NESTOR vacuum potential is not returned by the solver, so the wout potsin/xmpot/xnpot and *_sur variables are written as netCDF fill (see vmex.core.wout.wout_from_state()).
An NITER-exhausted free-boundary run still writes the wout (VMEC2000 behavior) and exits with ier_flag = 2 (MORE ITERATIONS REQUIRED).

vmex.core.cli.case_from_input(path: Path) → str¶: Case name from an input path (input.solovev -> solovev).

vmex.core.cli.resolve_wout_path(*, input_path: Path, outdir: Path | None) → Path¶: Default wout_<case>.nc path for an input file.

vmex.core.cli.build_parser() → ArgumentParser¶: Build the command-line parser for the vmex executable (vmec alias).

vmex.core.cli.main(argv: list[str] | None = None) → int¶: Run the vmec command-line entry point (zero-crash).

API reference¶

Inputs and profiles¶

Units¶

Spectral representation and physics kernels¶

VMEC2000 counterparts¶

Physics/numerics conventions (parity-critical)¶

VMEC2000 counterparts¶

Structure¶

VMEC2000 counterparts¶

Radial conventions (parity-critical)¶

VMEC2000 counterparts¶

VMEC2000 counterparts¶

Equations (Hirshman & Whitson 1983; VMEC2000 forces.f)¶

VMEC2000 counterparts¶

Conventions (parity-critical)¶

Solver¶

VMEC2000 counterparts¶

Representation conventions (parity-critical)¶

VMEC2000 counterparts¶

Equations (VMEC2000 notation)¶

Index conventions (parity-critical)¶

VMEC2000 counterpart¶

Modern (JAX) approach¶

VMEC2000 counterparts¶

Execution lanes (§5.3)¶

Structural executable reuse (2026-07-09, Phase 2 item (1))¶

Free boundary¶

Reuse, not re-implementation¶

ESSOS compatibility (essos.mgrid.MGrid, PR#33)¶

Differentiation and optimization¶

Residual formulation (documented choice)¶

Degrees of freedom / boundary handling¶

Gradient checking solver-sensitive metrics¶

Zero-crash typed errors through the callback¶

Parameter map¶

Helicity conventions (match legacy/simsopt exactly)¶

Gradient modes¶

Physics objectives¶

Scope notes¶

Scope notes¶

Scope notes¶

Outputs¶

VMEC2000 counterparts¶

Public API¶

Public API¶

Straight-axis mirrors¶

Errors and CLI¶

ESSOS compatibility (`essos.mgrid.MGrid`, PR#33)¶