Basic API

The API you use daily: the lazily imported vmex top-level exports (import vmex as vj) plus the three modules behind gradient-based work. Everything else — the per-module solver internals and the mirror lane — is in Advanced API.

Top-level package

Every name in vmex.__all__ is listed below by group, each linked to the module that documents it; this page is an index, not a second copy of every docstring.

vmex: a JAX implementation of VMEC2000 for fixed and free-boundary equilibria.

Every name below is a lazy attribute of the top-level package (import vmex as vj): the owning module is imported on first access, and vmex.__all__ lists exactly these names plus __version__. Each entry links to the module that documents it.

Inputs and run controls

Solvers

Outputs and scaling

High-order reconstruction and strong-force certificate

Force-balance polishing

Optimization

Post-processing and plotting

External fields

Errors, diagnostics, and modules

The vmex console entry point (vmec is an alias) lives in vmex.core.cli.

Inputs

VMEC input handling: the &INDATA namelist and structured 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 INDATA content this code base actually consumes, with VMEC2000 defaults. Parsing is host-side NumPy code (nothing here needs JAX). Controls which would change the mathematical problem or iteration contract but are not implemented are rejected by UnsupportedInputModeError; they are never silently converted into an ordinary fixed-/free-boundary solve.

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: the explicit legacy NS_ARRAY(1)=0 form expands to [max(3, NSIN), 31]; 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).

exception vmex.core.input.UnsupportedInputModeError(code: str, control: str, reason: str)

An input requests semantics which VMEX does not implement.

code is a stable, value-free diagnostic code suitable for the privacy-preserving input checker. control names the INDATA/JSON control without echoing its value or any equilibrium data.

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, lforbal: bool = False, lmove_axis: bool = True, lfull3d1out: bool = False, aphi: Any = None, phiedge: float = 1.0, nstep: int = 10, time_slice: float = 0.0, 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, restart_wout: str = '')

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)

lforbal: bool = False

replace m=1,n=0 R/Z forces by average force balance

lmove_axis: bool = True

improve the axis when the first force sum is > 1e2

lfull3d1out: bool = False

request the full threed1 output (J. Geiger)

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

time_slice: float = 0.0

informational value in the VMEC run header

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)

restart_wout: str = ''

wout path to seed the solve from (’’ = cold)

change_resolution(*, mpol: int, ntor: int, ntheta: int | None = None, nzeta: int | None = None) → VmecInput

Return a copy at the requested Fourier and real-space resolution.

Fourier and axis coefficients present at both resolutions are copied; newly added modes are zero. ntheta and nzeta keep their current values when omitted, including 0 for VMEC’s automatic grid choice. This method applies no optimization policy: callers choose every resolution explicitly.

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 structured 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 structured 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. The VMEC++ free_boundary_method="nestor" spelling is accepted. Other free-boundary methods and unknown keys fail explicitly instead of being silently 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.

DESC input bridge; WOUT is produced by the ordinary VMEX solve.

vmex.core.desc.is_desc_file(path: Path) → bool

Recognize DESC outputs or native text decks without importing DESC.

vmex.core.desc.write_desc_input(source: Path, outdir: Path | None = None, *, tolerance: float = 0.01) → Path

Read a DESC equilibrium without DESC and write a compact VMEC input.

The tolerance bounds discarded boundary position and angular derivatives, not the solved magnetic field. Zero retains all nonzero boundary modes.

Run directives (!@VMEX comment lines and the JSON _vmex section) are execution metadata and never become VmecInput fields; they are parsed and resolved here. The precedence rule is stated with the polishing options in Advanced API.

VMEX execution metadata carried outside the VMEC physics schema.

VMEX-only run controls travel in places every legacy reader ignores: comment directives in &INDATA text and a reserved _vmex section in structured JSON. VMEC2000 treats the directive lines as comments and solves the ordinary input; a VMEC++-compatible JSON reader sees the physics schema once _vmex is removed. VmecInput stays a pure physics object — these options control execution, not equilibrium physics, so they never become dataclass fields on it.

Two directive spellings are accepted, both VMEC-safe comments:

!@VMEX POLISH = AUTO
!@VMEX POLISH_TOL = 1.0E-8
!@VMEX POLISH_FAIL = ERROR
!@VMEX POLISH_DEGREE = 5
!@VMEX POLISH_MAX_ITER = 40
!@VMEX POLISH_SPANS = 16
!@VMEX POLISH_BUDGET = 3600

and the original single-flag form from the polishing integration:

! VMEX: POLISH_FORCE_BALANCE = .TRUE.

Precedence, resolved by resolve_run_options(), is exactly:

CLI option > explicit Python keyword > file directive > package default.
class vmex.core.run_options.InputRequest(input: VmecInput, options: RunOptions, source: Path)

One parsed input file: the physics and how to run it.

The result of read_input_request(), which is the only place VMEX builds one. It exists because the two halves of a deck have different lifetimes: the physics goes on to the solver unchanged, while the execution options are merged with CLI flags and Python keywords by resolve_run_options() before anything runs.

input

The parsed physics deck, a VmecInput. It carries no execution metadata: !@VMEX directive lines are Fortran comments to the &INDATA parser, and the JSON _vmex section is discarded by VmecInput.from_json_text.

Type:

VmecInput

options

The RunOptions read from those same directives or from the JSON _vmex section. A deck with no directives yields RunOptions(), i.e. every package default.

Type:

RunOptions

source

The path the deck was read from, kept so a caller can name outputs or resolve paths relative to the deck. VMEX itself only records it.

Type:

Path

class vmex.core.run_options.RunOptions(polish: bool | str = False, polish_tol: float | None = None, polish_fail: str = 'error', polish_degree: int | None = None, polish_max_iter: int | None = None, polish_spans: int | None = None, polish_budget: float | None = None)

How to execute a solve; never what equilibrium to solve.

polish is False, True, or "auto" (polish only when the legacy solve converged and the physics is in the supported set). polish_tol/polish_degree/polish_max_iter/polish_spans override the matching PolishConfig fields (tolerance, radial_degree, max_nonlinear_iterations, radial_spans) when set; only knobs that exist on the driver config are exposed here. polish_fail maps onto the driver’s fail policy: "error" raises, "fallback" returns the unpolished state silently, "warn" returns it with a RuntimeWarning. polish_budget is the wall-clock ceiling in seconds that POLISH = AUTO will commit to (auto_budget_seconds); it raises or lowers the cost at which AUTO declines to polish and has no effect on POLISH = ON, which never consults it.

vmex.core.run_options.format_indata_directives(options: RunOptions) → str

Serialize non-default options as canonical directive lines.

Only the !@VMEX KEY = VALUE spelling is written; the legacy ! VMEX: POLISH_FORCE_BALANCE form is read but never emitted. Round trips through parse_indata_run_options().

Parameters:

options – The options to serialize. A field equal to its RunOptions default (polish, polish_fail) or left at None (the four numeric overrides) is omitted, so an unmodified RunOptions() produces nothing.

Returns:

  • The directive lines as one string terminated by a newline, or the empty

  • string when every option is at its default. Prepend it to &INDATA

  • text (VMEC2000 reads the lines as comments.)

vmex.core.run_options.parse_indata_run_options(text: str) → RunOptions

Read every VMEX directive from raw &INDATA text.

Both spellings are read; a repeated key must repeat the same value, and a conflicting repetition is an error rather than a silent last-one-wins. Ordinary Fortran comments and quoted ! characters are unaffected: the directive patterns anchor on comment lines only, and VMEC2000 discards those lines entirely.

vmex.core.run_options.read_input_request(path: str | Path) → InputRequest

Read one input file into physics plus execution options.

VmecInput.from_file remains physics-only; this is the entry point the CLI and solve_file() share.

The format is chosen from the content, not only the name: a .json suffix or a first non-blank character of { selects the structured JSON reader, anything else the Fortran &INDATA reader. The file is read once as UTF-8 and both halves come from that same text, so the physics and the directives can never disagree about which revision was parsed.

Parameters:

path – Filesystem path to an &INDATA deck or a VMEC++-compatible JSON input.

Returns:

vmex.core.run_options.resolve_run_options(file_options: RunOptions | None, *, polish: bool | str | None = None, polish_tol: float | None = None, polish_fail: str | None = None, polish_degree: int | None = None, polish_max_iter: int | None = None, polish_spans: int | None = None, polish_budget: float | None = None) → tuple[RunOptions, dict[str, str]]

Apply the documented precedence and record where each value came from.

Explicit Python keywords override the file; the CLI passes its flags through the same keywords, so CLI > Python > file > default reduces to one merge. The source map ("python", "file", "default" per field) goes into the run report so a surprising activation is traceable.

Parameters:

  • file_options – Options parsed from the deck, normally InputRequest.options. None is treated as RunOptions().

  • polish – Explicit overrides, each with the meaning of the matching RunOptions field. None means “not specified” and leaves the file or default value in place, so an override cannot be used to reset a field back to None.

  • polish_tol – Explicit overrides, each with the meaning of the matching RunOptions field. None means “not specified” and leaves the file or default value in place, so an override cannot be used to reset a field back to None.

  • polish_fail – Explicit overrides, each with the meaning of the matching RunOptions field. None means “not specified” and leaves the file or default value in place, so an override cannot be used to reset a field back to None.

  • polish_degree – Explicit overrides, each with the meaning of the matching RunOptions field. None means “not specified” and leaves the file or default value in place, so an override cannot be used to reset a field back to None.

  • polish_max_iter – Explicit overrides, each with the meaning of the matching RunOptions field. None means “not specified” and leaves the file or default value in place, so an override cannot be used to reset a field back to None.

  • polish_spans – Explicit overrides, each with the meaning of the matching RunOptions field. None means “not specified” and leaves the file or default value in place, so an override cannot be used to reset a field back to None.

Returns:

  • ``(options, source)`` (the merged RunOptions, and a mapping)

  • from every RunOptions field name to "python", "file",

  • or "default" saying which layer supplied it. A field is credited

  • to "file" only when the deck’s value differs from the package

  • default.

vmex.core.run_options.strip_vmex_json(data: Mapping[str, Any]) → tuple[dict[str, Any], RunOptions]

Split a structured-JSON mapping into physics data and run options.

The physics schema remains VMEC++ compatible after _vmex is removed. Unknown _vmex keys fail explicitly — a typo must not silently run with defaults.

Differentiation and optimization

A converged vmex.core.optimize.Equilibrium exposes equilibrium.solution (the spectral equilibrium arrays) and equilibrium.solver_context (read-only grids, profiles, and constants). The shorter state and runtime attribute names remain compatible with existing code; runtime never means elapsed wall-clock time.

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

exactly: at the root the gradient needs only the multiplier of the raw force residual, whose Jacobian is block tridiagonal in radius, so one transposed block-Thomas solve refined once against an independent pullback, staged as one reusable per-config executable (_adjoint_block_core), returns g_p - mu^T dF_raw/dp with one more VJP — O(1) memory in the forward iteration count. Recycling GCROT(m, k) on the preconditioned transpose (_adjoint_gcrot_core) is the host fallback when that certificate fails. The adjoint solve executes host-eagerly on the caller’s thread and is bound to the config’s carried device on its own (see the “Adjoint execution site” section note); its convergence is enforced — a missed certificate raises the typed AdjointSolveError instead of silently returning a plausible-but-wrong gradient — and setting VMEX_ADJOINT_DEBUG=1 prints per-stage device/norm lines for hardware placement triage.

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.

Anchoring the fixed point

The theorem applies at a root of F, and the host solver does not stop at one: ftol gates the sum of SQUARES of the force, so a solve it reports converged still returns with |F| ~ sqrt(ftol) (2.7e-07 at ftol = 1e-12 on basic_non_stellsym_simsopt). Where dF/dz carries a small singular value — the lasym m=1 families, 1.5e-04 — that residual is a 1.8e-03 displacement of the state, far enough that a solver-sensitive metric read there is not the one whose derivative the adjoint computes. The host callback therefore Newton-refines the state onto the root (ImplicitConfig.refine_tol, _refined_state) before any lane reads it, so the value, the objective’s cotangent and the linearization all sit at the same point. Both must move together: with only the linearization refined the two errors stop cancelling (d(sum DMerc)/d(RBS(1,1)) against the frozen-path FD: rel 4.2e-03 at the host state, 5.7e-03 with the linearization alone refined, 5.4e-07 with both).

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 fixed-boundary dof mask is constructed directly from those mode-table invariants; _dof_mask() remains the independent structural-zero oracle and handles coupled free-boundary maps.

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 solver-sensitive metrics (iota at ncurr=1, mirror ratio, magnetic well, Boozer/QI residual) a naive re-solve FD is not a valid reference — it can sign-flip; use frozen_path_directional_fd(), which reproduces the adjoint to solver accuracy (full rationale on that function; tests/test_implicit_grad.py).

Strict and optimization-safe callback lanes

jax.pure_callback converts any host exception into an opaque JaxRuntimeError that embeds the whole host traceback and loses the typed exception (__cause__ is None) — breaking the vmex.core.errors zero-crash taxonomy. The strict diagnostic lane uses the relay: _host_solve_and_mask catches any VmecError, stashes it in the single-slot _HOST_ERROR (host callbacks are serialized per solve) and re-raises a SHORT sentinel RuntimeError; _callback_solve pops the slot and re-raises the ORIGINAL typed exception from None. im.run / solve_implicit() therefore fail with a short typed error. Under jax.jit the sentinel surfaces at the jit boundary instead (where the optimize.least_squares zero-crash penalty lanes catch it).

Optimization uses solve_implicit_status() instead. A typed equilibrium failure returns a fixed-shape fallback state and status code, after which the objective selects a finite differentiable penalty. Unexpected programming errors still propagate rather than masquerading as rejected trial points. An invalid initial point is rejected by a strict host preflight before an optimizer starts.

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, ac_aux_f: 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; ac_aux_f contains optimizable current-spline knot values while the knot positions remain static in the input; phiedge/pres_scale/curtor are 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, jacobian_adjoint_tol: float = 0.0001, jacobian_adjoint_maxiter: int = 10, adjoint_restart: int = 30, adjoint_maxiter: int = 300, adjoint_gcrot_m: int = 100, adjoint_gcrot_k: int = 20, refine_tol: float = 1e-10, max_fsq_ratio: float = 1000000.0, hot_restart: bool = False, device: Any = None)

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

jacobian_adjoint_tol: float = 0.0001

Tolerance for certifying the columns of an implicit residual Jacobian, kept separate from adjoint_tol because the two feed different consumers. A scalar gradient goes to a quasi-Newton method that accumulates curvature from it and wants every digit; a least-squares Jacobian only has to point a trust-region step, and the block factorization already backsolves every column before the certifier runs. Measured on the asymmetric quasi-axisymmetric stage: the certifier needs 542 iterations to reach 1e-6 and none to reach 1e-4, for a relative change in the Jacobian of 3.2e-5.

jacobian_adjoint_maxiter: int = 10

Restart budget for the Jacobian column certifier. It corrects a direct block-tridiagonal solve, so a handful of cycles either lands it or says the factorization has stopped being a good preconditioner at this iterate; spending the full adjoint_maxiter there buys nothing and costs half an hour per Jacobian on an asymmetric stage.

adjoint_gcrot_m: int = 100

reverse-adjoint GCROT(m, k) recycling solve (_adjoint_solve_gcrot): inner FGMRES cycle size m and k deflation directions. At high mode number plain restarted GMRES stalls short of adjoint_tol within its budget (the truncated cycle loses the small eigendirections); GCROT converges to the same lambda in fewer matvecs.

refine_tol: float = 1e-10

Newton-refine the host state onto the root of the frozen residual whenever |F(P(x*), p)| exceeds this (inf disables it). VMEC’s ftol gates the SUM OF SQUARES of the force, so a solve the host reports converged still leaves |F| ~ sqrt(ftol) — 2.7e-07 on basic_non_stellsym_simsopt at ftol = 1e-12, which the near-null m=1 direction there (singular value 1.5e-04) puts 1.8e-03 from the root. The implicit-function-theorem derivative holds only AT the root, so the lane returns the refined state and differentiates there: |F| drops to ~1e-14 and d(sum DMerc)/d(RBS(1,1)) from rel 4.2e-03 to 5e-07 against the frozen-path FD, for 14-26% of a forward solve and 9-14% of a value-and-gradient across the gradient decks.

max_fsq_ratio: float = 1000000.0

Largest (fsqr + fsqz + fsql) / ftol accepted for implicit differentiation when a trial exhausts its iteration budget.

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.

device: Any = None

the RESOLVED placement device (a jax.Device) or None. Carried in the static config so every stage of the operation — the pure_callback host solve (which runs on a runtime worker thread where the caller’s thread-local jax.default_device context does NOT apply), the cached runtime template, the custom-VJP backward and the multi-RHS pullback — re-enters the same device context on its own, without relying on an outer user-supplied context. Without this, a gradient evaluated outside such a context creates its constants on JAX’s default device and mixes devices with the committed state.

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.

class vmex.core.implicit.LinearResponseReport(residual_norm: Array, tolerance: Array, iterations: Array, converged: Array)

Per-right-hand-side residual certificate for an implicit response.

residual_norm: Any

Alias for field number 0

tolerance: Any

Alias for field number 1

iterations: Any

Alias for field number 2

converged: Any

Alias for field number 3

vmex.core.implicit.params_from_input(inp: VmecInput, *, device: Any = None) → ImplicitParams

Extract the differentiable parameters of an input as a pytree.

Omitted device (like explicit None) follows ordinary JAX placement. Pass device="auto" to request VMEX’s measured CPU preference for implicit Jacobians, or pass "cpu" / "gpu" or a jax.Device explicitly. High-level optimization entry points retain device="auto" as their default.

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.

Runs inside the config’s device context so freshly created constants (and the cached template) land on cfg.device even when the caller holds no outer jax.default_device context.

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, jacobian_adjoint_tol: float = 0.0001, jacobian_adjoint_maxiter: int = 10, adjoint_restart: int = 30, adjoint_maxiter: int = 300, adjoint_gcrot_m: int = 100, adjoint_gcrot_k: int = 20, refine_tol: float = 1e-10, max_fsq_ratio: float = 1000000.0, hot_restart: bool = False, device: Any = None) → ImplicitConfig

Build the static config; resolution is the (final-stage) grid.

device is the already-RESOLVED placement device (pass the result of vmex.core.device.resolve_implicit_device(), not a policy string) — see ImplicitConfig.device.

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_tangent_multi_rhs(params: ImplicitParams, cfg: ImplicitConfig, x_star: SpectralState, dof_mask: SpectralState, tangent_batch: ImplicitParams, *, probe_chunk_size: int | str = 1, response_chunk_size: int | str = 1) → tuple[SpectralState, LinearResponseReport]

State tangents for several parameter directions, factored once.

The pure-JAX raw-force Jacobian is exactly block tridiagonal in radius. One three-color assembly and SOLVAX factorization therefore initializes every right-hand side; a warm-started solve then certifies the ordinary preconditioned residual against 10 * cfg.adjoint_tol * ||rhs||. The two chunk sizes independently bound probe assembly and response solves; each is a positive int (default 1 — minimum memory, unchanged behavior) or opt-in "auto", which sizes the chunk from measured available memory on the config’s placement device divided by the per-column workspace computed exactly from the operand avals (see measured_chunk_size() for the memory model). The ordinary implicit reverse rule remains the default for solve_implicit().

vmex.core.implicit.implicit_state_pullback_multi_rhs(params: ImplicitParams, cfg: ImplicitConfig, x_star: SpectralState, dof_mask: SpectralState, gbar_batch: SpectralState, *, solver: str = 'gcrot', probe_chunk_size: int | str = 1, response_chunk_size: int | str = 1) → 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. solver="gcrot" (default) solves every row with preconditioned Krylov. solver="block" factors the raw nearest-neighbor radial Jacobian once and solves every row exactly with the transposed factors, refined once and certified like the scalar rule. Runs inside the config’s device context (see _solve_implicit_bwd); the two chunk sizes bound probe assembly and right-hand-side solves independently — each a positive int (default 1, unchanged behavior) or opt-in "auto", sized from measured available memory and the exact per-column operand bytes (see measured_chunk_size()).

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, refine_tol: float = 1e-10, device: Any = None) → ImplicitSolution

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

params defaults to params_from_input(); pass a perturbed / traced ImplicitParams to differentiate — and pass the SAME device to both calls (under jax.grad the parameters are tracers, which carry no placement to infer):

inp = VmecInput.from_file("input.solovev")
gpu = jax.devices("gpu")[1]
with device_scope(gpu):
    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).

device accepts "cpu", "gpu" or a jax.Device; omitted device and None leave placement to JAX. Use device_scope() around parameter creation and differentiation for a non-default accelerator. "auto": when params is supplied with CONCRETE arrays all committed to one device, that placement is preserved and used for the whole operation; otherwise (no params, tracers under jax.grad, or mixed placement) the default CPU preference of this launch-bound path applies — pass device= explicitly when differentiating on an accelerator. An explicit hardware device moves a supplied parameter pytree consistently.

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.measured_chunk_size(dim: int, per_column_bytes: int, *, device: Any = None, memory_fraction: float = 0.5) → int

Largest chunk whose per-column workspace fits measured free memory.

DESC sizes its objective jac_chunk_size from measured available memory divided by a hand-fit per-column constant (desc/objectives/objective_funs.py, desc/__init__.py); vmex knows the response operands’ avals exactly, so per_column_bytes is computed from the state pytree shapes instead. Memory model: a chunk holds chunk columns concurrently, each pinning per_column_bytes of workspace, and may claim memory_fraction of the measured available bytes, so chunk = memory_fraction * available // per_column_bytes clamped to [1, dim] (solvax.auto_chunk_size()’s explicit-budget regime). With no measurement available the same helper’s square-root heuristic bounds the chunk instead.

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: this differential-volume sum is pinned by ImplicitSolution.volume and the FD-cached gradient tables of tests/test_implicit_grad.py — keep it as is. 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: this shoelace-on-a-fresh-grid variant is pinned by ImplicitSolution.aspect and the FD-cached solovev gradient table of tests/test_implicit_grad.py — keep it as is. 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_axis(state: SpectralState, rt: SolverRuntime) → Any

On-axis rotational transform iotaf[0] of iota_profile() (differentiable).

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

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

Naming note: 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).

Naming note: 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.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.

The Newton steps go through the recycling GCROT solve rather than plain restarted GMRES: from a warm start already at the root the step’s RHS lies along the smallest eigendirection, where a truncated GMRES cycle stagnates. Measured on basic_non_stellsym_simsopt, that stall froze the -h branch at |F| = 2.2e-08 — 4e-04 from the root along a singular direction of dF/dz — and moved this FD by 35%.

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 — that stall is exactly what they catch).

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:

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 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 instead of one full equilibrium solve per dof (warm cost ~1.5 hot equilibrium solves independent of the dof count, vs one hot solve per dof for 2-point FD). In implicit mode every objective term must be a traceable function of (SpectralState, SolverRuntime); terms exposing a residuals_state method (QuasisymmetryRatioResidual) contribute their full pointwise residual vector (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; under jac="implicit" use d_merc_state() / mercier_stability_residual(), jdotb_residual(), and l_grad_b_state() instead. The implicit parameter map supports lasym via the four RBC/ZBS/RBS/ZBC boundary families and a traceable readin.f delta rotation (FD-validated); the traceable QS-ratio term follows, using the stored full poloidal grid for lasym states instead of mirroring VMEC’s reduced [0, pi] grid.

class vmex.core.optimize.VmecProblem(*args: Any, input_from_x: Callable[[Any], Any], x_from_input: Callable[[Any], Any], equilibrium_from_x: Callable[[...], Any] | None = None, boundary_from_x: Callable[[Any], Any] | None = None, **kwargs: Any)

A FunctionProblem backed by a VMEX equilibrium solve.

Adds the maps between the optimizer’s decision vector and VMEC objects: the input deck, the converged equilibrium, and the boundary coefficient arrays. It keeps the same optimizer contract as its base class, so the same methods go to SciPy, JAXopt, or Optax unchanged, and it enriches evaluate() with the solve and adjoint status of the underlying equilibrium.

Build one with from_tuples(), from_loss(), or from_input() rather than calling this constructor: they route through vmex.core.optimize.make_problem, which is what assembles the four callables below along with the objective, the derivative lane, the degree-of-freedom names, and the metadata.

Parameters:

  • *args – Forwarded positionally to FunctionProblem; in practice the decision vector x0.

  • **kwargs – Forwarded to FunctionProblem: the objective callables, names, bounds, scales, and metadata.

  • input_from_x – Required. x -> VmecInput: a new input deck carrying the boundary coefficients — and the current degrees of freedom, when the problem parameterizes them — of this decision vector. Nothing is solved.

  • x_from_input – Required. The inverse, VmecInput -> array: the decision vector that reproduces a given deck, which is the normal starting point of a continuation stage. x_from_input() rejects a result whose shape differs from x0.

  • equilibrium_from_x – Optional x -> Equilibrium, the converged equilibrium at x. Implicit problems return the accepted state the objective already computed rather than cold-solving the boundary again, which matters for strongly shaped boundaries whose cold axis guess can produce a sign-changing initial Jacobian. None makes equilibrium_from_x() raise AttributeError. A callable that accepts a newton_iterations keyword receives it only when the caller asks for something other than the default 10, so a closure without that keyword still works.

  • boundary_from_x – Optional x -> tuple of arrays, the traceable boundary coefficients: (rbc, zbs) for a stellarator-symmetric input and (rbc, zbs, rbs, zbc) when lasym. Each is a full dense INDATA-layout array of shape (2 * ntor + 1, mpol) indexed [n + ntor, m] in metres — not the trimmed decision vector — and is a JAX array, so it composes with coil or surface objectives under jax.grad(). None makes boundary_from_x() raise AttributeError.

classmethod from_tuples(inp: Any, objective_terms: Sequence[tuple[Callable[[...], Any], Any, float]], **kwargs: Any) → VmecProblem

Build a VMEC least-squares problem from weighted objective tuples.

The README entry point. Each tuple is (function, target, weight) and contributes one or more rows to a single residual vector; the rows of all terms are concatenated in the order given and the scalar cost is 0.5 * r @ r. Named row ranges are recorded in the problem metadata, which is what lets OptimizationMonitor report per-term costs without re-solving anything.

Parameters:

  • inp – The starting VmecInput. Its boundary supplies the initial decision vector and its resolution and profiles are held fixed apart from the parameterized degrees of freedom.

  • objective_terms –

    The (function, target, weight) triples.

    function is normally a traceable function(state, runtime) -> scalar or vector, evaluated on the converged equilibrium — the two arguments are also spelled (equilibrium_state, solver_context). An objective object exposing a residuals_state method may be passed instead (whole instance or bound method), in which case its full pointwise residual vector becomes this term’s rows. Under derivative_method="finite_difference" a one-argument host callable taking the whole Equilibrium is accepted too; the implicit lane rejects it, since it cannot be traced.

    target is the value the term is driven toward, coerced with float(), so it must be scalar even when function returns a vector — the same target is then subtracted from every row.

    weight is a non-negative scalar, or a one-dimensional array with one entry per residual row of that term. Under the default weight_semantics="cost" it multiplies the squared cost, so the row is sqrt(weight) * (function - target); with weight_semantics="residual" the row is weight * (function - target) and a negative entry is then allowed.

  • **kwargs – Passed through to vmex.core.optimize.make_problem: which boundary modes vary (max_mode, vary_major_radius, current_dofs), the derivative lane (derivative_method, implicit_jacobian_method, jacobian_batch_size), the forward solve controls, the variable scaling (use_ess, ess_alpha, bounds), and weight_semantics.

Returns:

  • A VmecProblem whose residual/residual_jac pair,

  • x0, and scales are ready for

  • scipy.optimize.least_squares(). A non-finite or empty

  • residual at the initial point raises FloatingPointError

  • rather than starting an optimization that cannot recover.

classmethod from_loss(inp: Any, loss: Callable[[...], Any], **kwargs: Any) → VmecProblem

Build a VMEC scalar problem from a traceable state/runtime loss.

Parameters:

  • inp – The starting VmecInput, as for from_tuples().

  • loss – loss(state, runtime) -> scalar, evaluated on the converged equilibrium and already carrying its own weights. It must return a single value: a vector-valued objective belongs in from_tuples(), or must be reduced here explicitly. Unlike an objective_terms entry it is used exactly as written — an object’s residuals_state is never substituted for it.

  • **kwargs – As for from_tuples().

Returns:

  • A VmecProblem exposing only the scalar lane —

  • fun(), grad(), and

  • their traceable counterparts — for a gradient optimizer such as

  • BFGS, L-BFGS-B, or Adam. It provides no residual or Jacobian.

classmethod from_input(inp: Any, **kwargs: Any) → VmecProblem

Parameterize an input for field VJPs without defining an objective.

Builds the same machinery as from_loss() around an identically zero loss, so there is nothing to minimize. Use it when what you want is the parameterization itself: the decision vector and its names, input_from_x() and boundary_from_x(), and the differentiable interior_field() and exterior_field() with exact VJPs in these degrees of freedom. inp and **kwargs are as for from_tuples().

input_from_x(x: Any) → Any

Return a new VmecInput containing decision vector x.

x_from_input(inp: Any) → ndarray

Return this problem’s decision vector for inp.

This is the inverse of input_from_x() for the boundary and any optional current degrees of freedom selected when the problem was constructed. It is the normal continuation-stage starting vector.

equilibrium_from_x(x: Any, *, newton_iterations: int = 10) → Any

Return the converged equilibrium evaluated at x.

Implicit problems reuse the accepted optimizer state instead of cold-solving the optimized boundary again. This matters for strongly shaped boundaries whose cold magnetic-axis guess may have a sign-changing initial Jacobian.

boundary_from_x(x: Any) → Any

Return traceable boundary coefficient arrays for decision vector x.

jax_objective_from_state(x: Any, extra_costs: Callable[[Any, Any], Any], *, n_extra_terms: int) → tuple[Any, tuple[Any, Any]]

Combine the VMEX least-squares cost with state-dependent costs.

extra_costs(state, runtime) returns one already-weighted scalar cost per added objective term. The auxiliary result contains the VMEX residual rows and those added costs, ready to pass as auxiliary data to jax.value_and_grad(). Failed equilibrium trials receive the same smooth finite rejection cost as the base problem, so driver scripts do not need their own accepted/rejected branches.

jax_extra_costs_from_state(x: Any, extra_costs: Callable[[Any, Any], Any], *, n_extra_terms: int) → tuple[Any, Any]

Evaluate additive state-dependent costs only at valid VMEC trials.

This is the split-compilation counterpart of jax_objective_from_state(). It returns zero extra cost at a rejected trial, leaving the base problem to supply its certified rejection wall. Splitting a large virtual-casing or coil graph from the VMEC objective substantially lowers peak XLA compilation memory.

jax_quantity_from_state(x: Any, quantity: Callable[[Any, Any], Any]) → tuple[Any, Any]

Evaluate a differentiable floating-point quantity and solve status.

quantity(state, runtime) may return any fixed-shape JAX array. Rejected equilibrium trials return an array of NaNs with that shape, so an invalid state cannot look like a usable diagnostic. The scalar status is zero only for an accepted equilibrium.

exterior_field(x: Any, *, external_field: Any | None = None, external_parameters: Any | None = None, external_field_from_parameters: Callable[[Any], Any] | None = None, external_dof_names: tuple[str, ...] = (), nphi: int = 32, ntheta: int = 32, digits: int = 6, levels: tuple[tuple[int, int], ...] | None = None, chunk_size: int | str = 'auto', target_chunk_size: int | str = 'auto') → Any

Return the exterior field and exact VJPs in this problem’s DOFs.

Query points must lie outside the last closed flux surface and away from coil filaments. The returned field follows the stored-point API: field.set_points(xyz); field.B(); field.B_vjp(cotangent). Set the source and target chunk sizes only to cap virtual-casing memory; "auto" is the tuned default.

interior_field(x: Any, *, newton_iterations: int = 10) → Any

Return the interior field and exact VJPs in this problem’s DOFs.

surface_field_values(x: Any, quantity: str, *, external_field: Any | None = None, nphi: int = 32, ntheta: int = 32, digits: int = 4, precision: Any | None = None) → Any

Return |B| or B.n/B on a trial boundary for plotting.

B.n/B is evaluated on the exterior side using the supplied coil or MGRID field plus the plasma-current virtual-casing field. This helper keeps optional movie coloring out of optimization driver code; it is not used by the objective or optimizer.

evaluate(x: Any, *, derivatives: bool = True) → Evaluation

Evaluate and attach VMEC solve/adjoint status diagnostics.

class vmex.core.optimize.FunctionProblem(x0: Any, *, fun: Callable[[ndarray], Any] | None = None, grad: Callable[[ndarray], Any] | None = None, value_and_grad: Callable[[ndarray], Any] | None = None, residual: Callable[[ndarray], Any] | None = None, residual_jac: Callable[[ndarray], Any] | None = None, residual_and_jac: Callable[[ndarray], Any] | None = None, jax_fun: Callable[[Any], Any] | None = None, jax_value_and_grad: Callable[[Any], tuple[Any, Any]] | None = None, jax_residual: Callable[[Any], Any] | None = None, jax_residual_jac: Callable[[Any], Any] | None = None, names: Sequence[str] | None = None, bounds: Any = None, scales: Any | None = None, metadata: Mapping[str, Any] | None = None, evaluation_progress: bool = False, report_interval: float = 10.0)

A decision vector plus optimizer-compatible objective callables.

Parameters are explicit and immutable from the caller’s perspective. Supplying combined value/gradient or residual/Jacobian functions enables a one-entry exact-key cache, so the common SciPy call sequence does not repeat expensive work. The cache is protected by a lock; JAX-native callables do not use host state and remain suitable for tracing.

This class deliberately does not provide solve(method=...). Pass its methods directly to the optimizer of choice.

At least one of fun, value_and_grad, residual, or residual_and_jac is required; every other callable is optional and the matching method raises AttributeError when its lane is absent. The host callables receive one contiguous float NumPy array and are free to be opaque; the jax_* callables receive whatever the caller traces and must stay traceable.

Two independent one-entry caches, each keyed on the exact bytes of x (shape, dtype, and contents — no tolerance), avoid repeating work across the split calls an optimizer makes at one iterate. The scalar cache is filled by value_and_grad() and so covers value_and_grad alone, fun together with grad, or residual_and_jac. The least-squares cache is filled by residual_and_jac() and so covers residual_and_jac, or residual together with residual_jac. Whether the cache actually pays depends on which callables were supplied: with residual_and_jac, SciPy’s separate fun(x) then jac(x) calls both route through it and the second is free, whereas separately supplied residual and residual_jac are each invoked directly and share nothing. Cached arrays are copied out, so a caller may mutate what it receives. The caches are guarded by a re-entrant lock; the jax_* lane touches no host state and stays safe to trace.

Parameters:

  • x0 – Initial decision vector. Copied to a float array; its size fixes the number of degrees of freedom and the expected Jacobian column count, and its shape is the shape gradients are reshaped to.

  • fun – x -> float, the scalar objective. Called directly by fun() without touching the cache.

  • grad – x -> array, the objective gradient. Used only in combination with fun; on its own it does not enable grad().

  • value_and_grad – x -> (float, array). The preferred scalar lane: it is one call for both quantities and it fills the scalar cache. The method that serves it is also reachable under SciPy’s fun_and_grad name.

  • residual – x -> array, the least-squares residual vector r(x). It is flattened, and defines the scalar objective 0.5 * r @ r when the problem has no scalar lane of its own (no fun, grad, or value_and_grad).

  • residual_jac – x -> array, the Jacobian dr_i/dx_j. It must have one column per decision variable; anything else raises ValueError.

  • residual_and_jac – x -> (array, array), the preferred least-squares lane: one call for both, filling the cache that residual() and residual_jac() then read. The Jacobian shape is checked exactly against (r.size, x0.size).

  • jax_fun – Traceable counterparts of the four callables above, returned unwrapped by the matching jax_* methods. They are never cached and never see host state, so they remain usable inside jax.jit() and jax.grad(). jax_fun falls back to the first element of jax_value_and_grad when it is not supplied.

  • jax_value_and_grad – Traceable counterparts of the four callables above, returned unwrapped by the matching jax_* methods. They are never cached and never see host state, so they remain usable inside jax.jit() and jax.grad(). jax_fun falls back to the first element of jax_value_and_grad when it is not supplied.

  • jax_residual – Traceable counterparts of the four callables above, returned unwrapped by the matching jax_* methods. They are never cached and never see host state, so they remain usable inside jax.jit() and jax.grad(). jax_fun falls back to the first element of jax_value_and_grad when it is not supplied.

  • jax_residual_jac – Traceable counterparts of the four callables above, returned unwrapped by the matching jax_* methods. They are never cached and never see host state, so they remain usable inside jax.jit() and jax.grad(). jax_fun falls back to the first element of jax_value_and_grad when it is not supplied.

  • names – One name per decision variable, in order, surfaced as dof_names for labelling output. The default is x[0], x[1], ...; a length other than x0.size raises ValueError.

  • bounds – Box constraints stored verbatim for the optimizer to consume — a SciPy Bounds, or a (lower, upper) pair. This class neither interprets nor enforces them.

  • scales – Positive finite per-variable scale factors with the shape of x0, defaulting to ones. Like bounds they are carried, not applied: pass them to the optimizer (SciPy’s x_scale). A non-finite or non-positive entry raises ValueError.

  • metadata – Free-form mapping copied onto the instance. VMEX-built problems use it to carry the named residual slices, the solver configuration, the mutable solve counters, and the traceable state accessors that VmecProblem reads.

  • evaluation_progress – Print an elapsed-time heartbeat around long evaluations. It stays silent until a call outlives the first interval, so fast calls print nothing. It wraps the standalone residual() and residual_jac() calls only, which is where a production deck spends minutes; the combined and scalar lanes are unaffected.

  • report_interval – Seconds between heartbeat lines. Must be positive.

property dof_names: tuple[str, ...]

Ordered names corresponding one-to-one with entries of a decision vector.

classmethod from_functions(x0: Any, **kwargs: Any) → FunctionProblem

Build a problem from user-supplied x-level callables.

value_and_grad(x: Any) → tuple[float, ndarray]

Return scalar value and gradient for SciPy jac=True.

fun_and_grad(x: Any) → tuple[float, ndarray]

Return scalar value and gradient for SciPy jac=True.

fun(x: Any) → float

Return the scalar objective value.

grad(x: Any) → ndarray

Return the scalar objective gradient.

residual_and_jac(x: Any) → tuple[ndarray, ndarray]

Return the residual vector and its Jacobian.

residual(x: Any) → ndarray

Return the residual vector.

residual_jac(x: Any) → ndarray

Return the residual Jacobian.

J(x: Any) → float

SIMSOPT-style alias for fun().

dJ(x: Any) → ndarray

SIMSOPT-style alias for grad().

jax_fun(x: Any) → Any

Return the traceable scalar objective.

jax_value_and_grad(x: Any) → tuple[Any, Any]

Return the traceable scalar value and gradient.

jax_residual(x: Any) → Any

Return the traceable residual vector.

jax_residual_jac(x: Any) → Any

Return the traceable residual Jacobian.

evaluate(x: Any, *, derivatives: bool = True) → Evaluation

Evaluate available scalar and residual quantities at x.

compile_residual_and_jacobian(x: Any | None = None, *, progress: bool = True, report_interval: float = 10.0, stream: Any = None) → Evaluation

Compile and cache the least-squares residual and Jacobian.

This call is optional: an optimizer compiles on its first evaluation if it is omitted. Calling it explicitly provides elapsed-time output during a potentially long first JAX compilation. Later calls at the same x use the normal one-entry problem cache.

compile_value_and_gradient(x: Any | None = None, *, progress: bool = True, report_interval: float = 10.0, stream: Any = None) → Evaluation

Compile and cache the scalar value and gradient.

This optional call makes the first JAX compilation visible before BFGS, L-BFGS-B, Adam, or another gradient optimizer starts.

class vmex.core.optimize.Evaluation(x: ndarray, value: float | None = None, gradient: ndarray | None = None, residual: ndarray | None = None, jacobian: ndarray | None = None, status: str = 'success', message: str = '', diagnostics: Mapping[str, ~typing.Any]=<factory>)

Values and diagnostics produced at one decision vector.

Fields that were not requested or are unavailable are None. status is a short machine-readable value such as "success" or "failed_solve"; message is intended for a human. Optimizers use the ordinary callable methods and do not need to understand this object.

Returned by FunctionProblem.evaluate() and the two compile_* helpers. It is a report, not a cache: nothing here is consulted by a later evaluation.

x

The decision vector the values were produced at, as an owned float copy with the shape of problem.x0.

Type:

numpy.ndarray

value

Scalar objective at x. For a residual-only problem this is the least-squares cost 0.5 * r @ r, matching SciPy’s OptimizeResult.cost. None when the problem exposes neither a scalar objective nor residuals.

Type:

float | None

gradient

Gradient of value with respect to x, reshaped to x’s shape. For a residual problem it is J.T @ r. None when derivatives=False was requested or no gradient lane exists.

Type:

numpy.ndarray | None

residual

The flattened residual vector r(x), or None for a scalar-only problem.

Type:

numpy.ndarray | None

jacobian

The residual Jacobian dr_i/dx_j with shape (residual.size, x.size), or None when derivatives were not requested or the problem provides no Jacobian.

Type:

numpy.ndarray | None

status

"success", or — from a VMEC-backed problem — "failed_solve" when the equilibrium solve at x raised, and "under_converged" when it returned but its force residual exceeds the threshold below which implicit derivatives are certified.

Type:

str

message

Human-readable explanation, empty on success; the solver exception’s text for "failed_solve".

Type:

str

diagnostics

Extra per-evaluation values. Empty for a plain FunctionProblem. A VmecProblem adds the cumulative failed_trials and derivative_fallbacks counters, a solve_stats mapping when the implicit lane recorded one, and — when the equilibrium at x could be materialised — the summed force residual fsq, its fsq_ratio to the solve tolerance, the configured max_fsq_ratio, and the boolean derivative_certified. A failed solve also carries exception_type.

Type:

Mapping[str, Any]

property success: bool

Whether the evaluation completed without a recoverable failure.

class vmex.core.optimize.EquilibriumReporter(*quantities: tuple[str, ~typing.Callable[[...], ~typing.Any], str], stream: ~typing.TextIO | None | object = <object object>, separator: str = ', ')

Print a compact set of scalar diagnostics for an equilibrium.

Each quantity is (label, callable, format_spec). Callables may use either the function(equilibrium) or function(state, runtime) convention used by VMEX objectives. Calling the reporter prints one line and returns the values by label, so scripts can also reuse a final metric.

Parameters:

  • *quantities – One (label, function, format_spec) triple per reported column; at least one is required and the labels must be unique. label names the column and keys the returned mapping. function is dispatched on its signature: a callable whose second positional parameter has no default is called as function(state, runtime) (the VMEX objective convention), every other callable as function(equilibrium). It must return exactly one scalar; any other size raises ValueError. format_spec is a format() specification applied to that float, for example ".6e".

  • stream – Where the report line is written. The default is sys.stdout; pass None to compute and return the values without printing.

  • separator – Text placed between the label = value fields of the printed line.

class vmex.core.optimize.OptimizationMonitor(problem: FunctionProblem | None = None, *, stream: TextIO | None | object = <object object>, print_every: int = 1, trace: bool = True)

Record and optionally print optimizer iterations and trials.

Pass the instance as a SciPy callback. SciPy invokes callbacks after an iteration, unlike objective functions which are also called for rejected line-search or trust-region trials. JAXopt, Optax, and custom loops can call record() with values they already computed. trace also prints one trial line per objective evaluation, so the rejected line-search trials between two accepted iterations – a full equilibrium solve each – are visible while the run is in progress.

The monitor never chooses steps or changes an optimizer. If problem is supplied, VMEX solve/failure counters are read without evaluating the objective again.

Parameters:

  • problem – Optional problem the optimizer is running on. It is used only to read metadata: the named residual slices that split cost into per-term costs, the cumulative equilibrium-solve and failed-trial counters, and — as a last resort, when a callback carries neither cost nor fun — one fun() call at the accepted iterate. With None those record fields stay None or empty and the monitor still records everything the callback provides.

  • stream – Where the per-iteration table is written. The default is sys.stdout; pass None to record silently and read records, history, or save() afterwards.

  • print_every – Print one row every print_every records (the first record is always printed, together with the header). Must be at least 1; recording is unaffected.

wrap_value_and_grad(function: Callable | Sequence[Callable], term_names: tuple[str, ...] | None = None, *, residual_slices: tuple[tuple[str, int, int], ...] = ()) → Callable

Adapt a JAX has_aux value/gradient pair for SciPy.

Each function(x) must return ((cost, terms), gradient). Pass a sequence to compile large additive physics components separately; their costs and gradients are summed without changing the optimizer contract. terms may be a mapping of labels to weighted scalar costs, or one compact vector paired with term_names. For a large residual graph, pass residual_slices and return (residual, extra_costs...); costs are reduced on the host to keep the compiled output small. The first evaluation is recorded as iteration zero; later evaluations are cached.

cache_evaluation(x: Any, cost: Any, gradient: Any, terms: Mapping[str, Any] | None = None) → tuple[float, ndarray]

Cache an already-computed objective pair for a SciPy callback.

This method does no differentiation and does not alter the objective. Driver scripts can show their explicit jax.value_and_grad calls, sum independently compiled physics components themselves, and use this one host conversion to avoid recomputing per-term costs when SciPy later reports an accepted iterate.

record(x: Any, *, cost: float, optimality: float | None = None, iteration: int | None = None, equilibrium_solves: int | None = None, rejected_trials: int | None = None, terms: Mapping[str, Any] | None = None, counters: Mapping[str, Any] | None = None) → OptimizationRecord

Append one already-computed accepted iterate and return its record.

Parameters:

  • x – The accepted decision vector; a float copy is appended to x_history.

  • cost – Total scalar objective at x. Nothing is recomputed.

  • optimality – First-order optimality measure, or None when unknown.

  • iteration – Iteration index; defaults to the number of records already held, and is advanced past the previous record when an optimizer restarts its own counter at a continuation stage.

  • equilibrium_solves – Cumulative solve and failed-trial counts. None (the default) reads them from the monitor’s problem, leaving the field None when no problem was supplied.

  • rejected_trials – Cumulative solve and failed-trial counts. None (the default) reads them from the monitor’s problem, leaving the field None when no problem was supplied.

  • terms – Per-term weighted costs by label. None splits the problem’s own residual over its named term slices instead, which may evaluate the residual once at x.

  • counters – Effort counters by name. None reads them from the monitor’s problem, empty when it carries none.

Return type:

The appended OptimizationRecord.

property history: dict[str, ndarray]

Return total and per-term cost histories as one mapping.

property x_history: tuple[ndarray, ...]

Copies of the accepted decision vectors, including iteration zero.

save(path: str | Path) → Path

Save the recorded iteration and cost columns as CSV.

plot(path: str | Path, *, title: str = 'Optimization objective terms') → Path

Write a compact log-scale total and per-term cost history plot.

movie(path: str | Path, object_factory: Callable[[ndarray], Any], **kwargs: Any) → Path

Animate accepted surface/coil iterates with one geometry callback.

movie_surface_coils(path: str | Path, object_factory: Callable[[ndarray], Any], *, x0: Any, scales: Any, surface_color: str | Callable | None = None, plasma_problem: Any | None = None, external_field: Callable[[Any], Any] | None = None, nphi: int = 32, ntheta: int = 32, digits: int = 4, precision: Any | None = None, **kwargs: Any) → Path

Animate normalized surface/coil iterates with optional field color.

object_factory receives physical variables x0 + scales*u. surface_color may be None, "absB", "B.n/B", or a callable (u, objects) -> values. For the two named field colors, provide plasma_problem; B.n/B additionally needs external_field(objects). These plotting helpers never enter the objective or gradient graph.

class vmex.core.optimize.OptimizationRecord(iteration: int, cost: float, reduction: float | None, optimality: float | None, equilibrium_solves: int | None, rejected_trials: int | None, terms: Mapping[str, float]=<factory>, counters: Mapping[str, float | None]=<factory>)

One optimizer callback, normally one accepted iteration.

Produced by OptimizationMonitor.record() and stored in OptimizationMonitor.records. Every field is a plain host value; a field is None when the optimizer callback did not supply it and the monitor could not derive it, never zero-as-unknown.

iteration

Iteration index of this accepted iterate. Taken from the SciPy result’s nit when present, otherwise the number of records already held. Optimizers restart their counter at every continuation stage, so a repeated or decreasing value is bumped to one past the previous record; the sequence in one monitor is always strictly increasing.

Type:

int

cost

Total scalar objective at this iterate. For a least-squares problem this is 0.5 * r @ r over the full residual vector, so it matches SciPy’s OptimizeResult.cost; for a scalar objective it is the objective value itself. Units are those of the weighted objective (dimensionless for the usual normalised VMEX terms).

Type:

float

reduction

previous cost - this cost, positive when the step improved the objective. None for the first record, which has no predecessor.

Type:

float | None

optimality

First-order optimality measure of the gradient. SciPy’s own optimality is used when the callback provides it; otherwise it is the infinity norm of a callback-supplied jac, or the Euclidean norm of the gradient cached by OptimizationMonitor.cache_evaluation(). None when no gradient information reached the monitor.

Type:

float | None

equilibrium_solves

Cumulative number of forward VMEC equilibrium solves performed for this problem’s implicit configuration, read from the solver’s own counters without re-evaluating the objective. None when the monitor was built without a problem or the problem carries no implicit configuration (for example a finite-difference or wout-only problem).

Type:

int | None

rejected_trials

Cumulative number of optimizer trial points whose equilibrium solve failed and was replaced by the smooth rejection cost. None under the same conditions as equilibrium_solves.

Type:

int | None

terms

Per-term weighted costs by label, 0.5 * r_k @ r_k over each named residual slice of the problem, so the values sum to cost for a pure least-squares problem. Empty when the problem exposes no term slices and the caller supplied none.

Type:

Mapping[str, float]

counters

Cumulative effort counters of the problem’s implicit configuration at this record: solves and descent iterations, refinement calls, steps and Krylov iterations, Jacobian calls, columns and certifier Krylov iterations, adjoint calls and Krylov iterations, and the host seconds of each part. An entry is None when the work ran inside a compiled program where it cannot be observed. Empty under the same conditions as equilibrium_solves.

Type:

Mapping[str, float | None]

vmex.core.optimize.make_problem(inp: VmecInput, *, objective_terms: Callable, float, ~typing.Any]] | None=None, loss: Callable | None = None, max_mode: int = 1, vary_major_radius: bool = False, x0: ndarray | None = None, current_dofs: int | None = None, derivative_method: str = 'implicit', fd_method: str = '3-point', fd_rel_step: float | None = None, workers: int | None = None, weight_semantics: str = 'cost', jacobian_batch_size: int | str | None = 1, implicit_jacobian_method: str = 'auto', adjoint_tol: float = 1e-06, jacobian_adjoint_tol: float = 0.0001, jacobian_adjoint_maxiter: int = 10, adjoint_maxiter: int = 300, max_fsq_ratio: float = 1000000.0, refine_tol: float = 1e-10, forward_ftol: float | None = None, forward_max_iterations: int | None = None, hot_restart: bool = True, warm_start: str | None = 'perturbation', use_ess: bool = True, ess_alpha: float = 1.2, evaluation_progress: bool = False, bounds: Any = None, device: Any = 'auto', solve_kwargs: dict | None = None, restart_from: Any = None, progress: bool = False, report_interval: float = 10.0, progress_stream: Any = None, problem_class: type[VmecProblem] = <class 'vmex.core.problem.VmecProblem'>) → VmecProblem

Build optimizer-neutral VMEC objective and derivative callables.

Exactly one of objective_terms and loss is required. With the default weight_semantics="cost", tuple weight w multiplies the squared cost, so the residual row is sqrt(w) * (f - target). Select weight_semantics="residual" when w should multiply the residual itself. loss must be a traceable (state, runtime) -> scalar callable.

derivative_method="implicit" computes exact derivatives of the converged fixed-boundary equilibrium by implicit differentiation and requires traceable objectives. "finite_difference" also accepts opaque host objectives and uses independent equilibrium re-solves; workers=None selects the automatic host-worker count. Users with complete x-level derivatives can use FunctionProblem.from_functions().

implicit_jacobian_method="auto" is the beginner-facing default: it selects one reverse adjoint for a scalar residual and an amortized block-tridiagonal factorization for vector residuals. Advanced choices are "block_tridiagonal", "forward_gmres", and "reverse_adjoint"; the names describe how the exact implicit Jacobian is assembled. An uncertified automatic response is recomputed by reverse adjoint. Forced host methods raise AdjointSolveError; transformable JAX methods use the same reverse fallback because Python exceptions cannot be raised reliably under jit.

jacobian_batch_size=1 is the default for QI/QS problems through max_mode=5: it minimizes cold compilation complexity and peak memory. "auto" batches response columns and improves warm throughput, so it is preferable for long same-shape continuation campaigns that amortize the larger first compilation. This public name maps to the compatibility drivers’ established jac_chunk_size implementation.

Set progress=True to report elapsed-time heartbeats while validating the seed equilibrium and building resolution-dependent solver data. VmecProblem.compile_residual_and_jacobian() or VmecProblem.compile_value_and_gradient() provides the same visibility for the first derivative evaluation after this factory returns.

The returned object contains no optimization algorithm. Pass VmecProblem.residual() / VmecProblem.residual_jac() to a nonlinear least-squares package, or VmecProblem.value_and_grad() to any scalar gradient optimizer.

forward_ftol and forward_max_iterations override the final VMEC solve stage for either derivative method. max_fsq_ratio controls how close an iteration-limited trial must be to that tolerance before VMEX differentiates it. The default accepts FSQ / forward_ftol <= 1e6; stricter studies can reduce it without changing VMEX internals.

adjoint_tol is a relative Krylov tolerance with a certified true residual check; adjoint_maxiter is the restart budget.

refine_tol is the frozen fixed-point residual required before implicit differentiation. The default preserves strict gradients; numpy.inf explicitly disables refinement for legacy comparisons.

restart_from seeds the first equilibrium from a previous WOUT, Equilibrium, or solver result. This is useful when a continuation stage changes mpol, ntor, or radial resolution: trial hot restarts then continue from the remapped converged state instead of a cold axis guess.

class vmex.core.optimize.Equilibrium(inp: VmecInput, state: SpectralState, runtime: SolverRuntime, result: SolveResult, field_factory: Callable[[], Any] | None = None, exterior_field_factory: Callable[[...], Any] | None = None)

A converged fixed-boundary equilibrium plus its evaluation contexts.

Objective callables in least_squares() receive one of these. The solution and solver_context are the clear public names for the solver-native state and runtime attributes. They feed the differentiable scalar targets; wout (built lazily, host NumPy) feeds wout-table objectives (QS ratio residual, Boozer-based QI residual).

property solution: SpectralState

Converged spectral equilibrium coefficients and force arrays.

property solver_context: SolverRuntime

Read-only grids, profiles, and constants used to evaluate the solution.

property wout: WoutData

Full wout dataset of this state (vmex.core.wout, cached).

exterior_field(**kwargs)

Return a field that can be queried outside the plasma surface.

property field

Pointwise magnetic field inside the plasma boundary.

set_points(points: Any) → Equilibrium

Store Cartesian points for pointwise field evaluation.

set_points_xyz(points: Any) → Equilibrium

Store Cartesian (x, y, z) points for field evaluation.

set_points_cyl(points: Any) → Equilibrium

Store cylindrical (R, phi, Z) points for field evaluation.

set_points_flux(points: Any) → Equilibrium

Store VMEC (s, theta, phi) points inside the plasma.

field_in_flux_coordinates()

Return the interior field in the (s, theta, phi) basis.

B(points: Any | None = None) → Any

Return Cartesian B inside the plasma.

absB(points: Any | None = None) → Any

Return |B| inside the plasma.

gradB(points: Any | None = None) → Any

Return the Cartesian field gradient inside the plasma.

gradgradB(points: Any | None = None) → Any

Return the second Cartesian derivative of B.

gradgradgradB(points: Any | None = None) → Any

Return the third Cartesian derivative of B.

B_vjp(vector: Any) → Any

Return the VJP of B() in the originating problem’s DOFs.

gradB_vjp(vector: Any) → Any

Return the VJP of gradB() in the problem’s DOFs.

gradgradB_vjp(vector: Any) → Any

Return the VJP of gradgradB() in the problem’s DOFs.

gradgradgradB_vjp(vector: Any) → Any

Return the VJP of gradgradgradB() in the problem’s DOFs.

vmex.core.optimize.solve_equilibrium(inp: VmecInput, *, initial_state: SpectralState | None = None, raise_on_max_iterations: bool = False, verbose: bool = False, forward_ftol: float | None = None, forward_max_iterations: int | None = None, polish_force_balance: bool | str = False, **solve_kwargs) → Equilibrium

Converge inp with the core multigrid solver -> Equilibrium.

verbose=True prints the VMEC iteration table, including the current iteration count and force residuals. 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(). forward_ftol and forward_max_iterations replace the input ladder’s final tolerance and iteration cap, using the same names as the optimization problem API. Force-balance polishing is strictly opt-in: the default polish_force_balance=False returns the ordinary converged VMEC state.

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.min_abs_iota(state: SpectralState, rt: SolverRuntime) → Any

Smallest |iota| over the half-mesh surfaces (axis excluded).

The transform floor that matters for a stellarator is the profile minimum, not its average: a mean target is satisfiable while a shear reversal or a current-carried profile leaves an interior surface near zero transform. The absolute value keeps the metric sign-free, since only the magnitude of the transform is physical here.

vmex.core.optimize.soft_min_abs_iota(state: SpectralState, rt: SolverRuntime, *, tau: float = 0.02) → Any

Smooth stand-in for min_abs_iota().

The hard minimum is differentiable except where two surfaces tie, which a least-squares step can sit on. This variant weights the profile by softmax(-|iota| / tau), so it stays inside [min, max] and reduces to the minimum as tau -> 0; tau carries the units of iota, so 0.02 resolves a minimum to a few percent. The log-sum-exp softmin is deliberately avoided: it sits tau log(ns) below the true minimum, which turns a floor on a non-negative quantity into a negative number.

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: 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]).

Naming note: 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.mirror_ratio(state: SpectralState, rt: SolverRuntime, *, s_index: int = -1) → Any

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

Note the convention: this is the |B| modulation depth on a surface, the standard QI optimization knob, not R_m = Bmax / Bmin. The two are related by R_m = (1 + m) / (1 - m). The open-mirror lane reports R_m proper — R_m,axis per leg and R_m,LCFS separately — through vmex.mirror.metrics.

|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.volume_average_beta(state: SpectralState, rt: SolverRuntime) → Any

Volume-averaged plasma beta, equal to the wout betatotal scalar.

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

simsopt’s Vmec.vacuum_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. Ported from legacy vmex.finite_beta.magnetic_well_from_vp.

vmex.core.optimize.elongation_profile(state: SpectralState, rt: SolverRuntime, *, ntheta: int | None = None, nphi: int | None = None) → Array

Boundary elongation at equally spaced toroidal cross-sections.

Each constant-geometric-toroidal-angle cross-section is reconstructed from the physical edge Fourier coefficients. Its area A and perimeter P define the equivalent ellipse, whose semi-axes are recovered with Ramanujan’s perimeter approximation. The returned ratio is a_major / a_minor at each toroidal sample, the same definition used by DESC’s elongation objective.

The default grids (at least 32 poloidal and 24 toroidal points per field period) resolve optimization modes through max_mode <= 5 without adding user-facing sampling knobs to ordinary scripts. Explicit positive ntheta and nphi values are available for convergence studies. The calculation is JAX-traceable and supports stellarator-asymmetric boundaries.

vmex.core.optimize.max_elongation(state: SpectralState, rt: SolverRuntime, *, ntheta: int | None = None, nphi: int | None = None) → Any

Maximum boundary elongation over one field period.

See elongation_profile(). A hard maximum is deliberate: it matches the usual engineering constraint and is differentiable except when two sampled cross-sections tie exactly.

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. Use traceable mercier_stability_residual() with jac="implicit". Both symmetry modes are supported; the lasym lane is validated per-term against live VMEC2000 output (tests/test_vmec2000_live.py).

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

Traceable VMEC DMerc profile on the full radial mesh.

Positive interior values indicate Mercier stability. This is a pure-JAX port of the symmetric jxbforce.f/mercier.f path used by vmex.core.nyquist.mercier_and_jxb(); it accepts a live converged (state, runtime) pair and supports jit, JVP and reverse-mode AD. The axis, first near-axis surface and edge retain VMEC’s zero/noisy output convention and should be excluded from objectives (normally [2:-1]).

Supports both symmetry modes: the lasym lane carries the sine-parity contributions through the jxbforce.f parity-split filter and the full-theta-interval surface integrals, validated per-term against live VMEC2000 lasym output (tests/test_vmec2000_live.py) and against an independent NumPy reconstruction from the wout Fourier tables (vmex.core.plotting._glasser_d_r_from_wout()).

vmex.core.optimize.mercier_stability_residual(state: SpectralState, rt: SolverRuntime, *, margin: float = 0.0, smoothing: float = 1e-06) → Any

Smooth Mercier-instability residual on DMerc[2:-1].

Positive DMerc is stable. For each validated interior surface this returns smoothing * softplus((margin - DMerc) / smoothing): it tends to max(margin - DMerc, 0) as smoothing tends to zero, while retaining a smooth gradient at the stability boundary. At finite smoothing it is strictly positive but exponentially close to zero on a sufficiently stable surface. Use target zero in vmex.core.optimize.least_squares(); margin > 0 requests a finite stability margin. The first two surfaces and edge are excluded.

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

Traceable VMEC jdotb = <J.B> profile in WOUT units.

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

Interior <J.B> profile for least-squares current objectives.

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

Return S = d(iota)/d(Phi) in the VMEC Mercier normalization.

vmex.core.optimize.glasser_d_r_state(state: SpectralState, rt: SolverRuntime, *, shear_epsilon: float = 0.0) → Any

Traceable Glasser–Greene–Johnson D_R profile.

Non-positive values satisfy the necessary local resistive-interchange stability condition on nonzero-shear surfaces. With the strict default, exact zero-shear entries are set to zero because the criterion is undefined there. A positive shear_epsilon replaces the denominator shear**2 by shear**2 + shear_epsilon**2 for smooth optimization; this regularization does not make zero-shear surfaces physically valid. Post-check mercier_shear_state() and require every target surface to satisfy abs(S) >> shear_epsilon before interpreting the result. As for DMerc, use only validated interior surfaces (normally [2:-1]) as optimization targets.

Supports both symmetry modes. The symmetric lane retains its DCON comparison; the lasym lane is validated through the VMEC2000-anchored DMerc parity plus the exact in-repo identity D_R = -DMerc + (H - S²/2)²/S² (no external lasym D_R oracle exists — the DCON comparison is symmetric-only).

vmex.core.optimize.glasser_stability_residual(state: SpectralState, rt: SolverRuntime, *, margin: float = 0.0, smoothing: float = 1e-06, shear_epsilon: float = 1e-08) → Any

Smooth resistive-interchange residual on D_R[2:-1].

Subject to the prerequisite DMerc > 0, stable surfaces require D_R <= 0. Combine this residual with mercier_stability_residual(); targeting it to zero penalizes D_R > -margin while retaining a smooth derivative. The nonzero default shear regularization makes this optimization helper finite on zero-shear seeds; use glasser_d_r_state() with its strict default for reporting.

vmex.core.optimize.trial_pressure_d_merc_state(state: SpectralState, rt: SolverRuntime, *, beta: float = 0.025, pressure_shape: Any = None) → Any

Mercier profile for a trial pressure on frozen equilibrium geometry.

This inexpensive AD-transparent proxy replaces only the explicit pressure drive in VMEC’s Mercier expression. It does not include the finite-beta geometry, current, or Shafranov-shift response; certify a candidate with a self-consistent finite-pressure equilibrium.

vmex.core.optimize.trial_pressure_glasser_d_r_state(state: SpectralState, rt: SolverRuntime, *, beta: float = 0.025, pressure_shape: Any = None, shear_epsilon: float = 0.0) → Any

Glasser D_R for the same frozen-equilibrium trial pressure proxy.

vmex.core.optimize.trial_pressure_mercier_stability_residual(state: SpectralState, rt: SolverRuntime, *, beta: float = 0.025, pressure_shape: Any = None, margin: float = 0.0, smoothing: float = 1e-06) → Any

Smooth instability residual for trial-pressure DMerc[2:-1].

vmex.core.optimize.trial_pressure_glasser_stability_residual(state: SpectralState, rt: SolverRuntime, *, beta: float = 0.025, pressure_shape: Any = None, margin: float = 0.0, smoothing: float = 1e-06, shear_epsilon: float = 1e-08) → Any

Smooth instability residual for trial-pressure D_R[2:-1].

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))`` (squared Frobenius norm of the Cartesian field-gradient tensor) — the Kappel/Landreman coil-complexity / compactness proxy. Evaluated from the wout tables of the converged state (bsupumnc/bsupvmnc Nyquist spectra, spectral angular derivatives of rmnc/zmns, native half/full-mesh radial finite differences, one-sided at the edge); 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; asymmetric inputs raise instead of silently dropping their sine/cosine partners. Accepts an Equilibrium or wout-like. 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(): identical convention and grid, with the wout coefficient tables rebuilt traceably from the state (_lgradb_state_tables(), the wrout.f Nyquist analysis as jnp einsums) and the same 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, differentiable almost everywhere, 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 smooth, report the hard minimum.

vmex.core.optimize.quasi_isodynamic_residual(*, bmnc_b, bmns_b=None, 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 (the default weights reproduce exactly the terms used by the established minimal-seed QI formulation):

  • 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).

Legacy port (quasi_isodynamic_residual_from_boozer_modes) with the unused aligned_profile_* / weighted_shuffle_* / shuffle_profile_nphi_out machinery removed. 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, bmns_b, xm_b, xn_b, iota_b, nfp, s_b} with the spectra shaped (nsurf, nmodes). bmns_b is zero for symmetric equilibria and contains the independent sine spectrum for lasym.

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, *, vary_major_radius: bool = False) → 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.boundary_arrays_from_x(inp: VmecInput, x, max_mode: int, *, vary_major_radius: bool = False) → tuple[Any, ...]

Traceable boundary coefficient arrays reconstructed from x.

The default returns (rbc, zbs); an asymmetric input additionally returns (rbs, zbc). vary_major_radius=True appends RBC(0,0) to the decision vector without introducing the identically-zero ZBS(0,0) direction.

vmex.core.optimize.pack_boundary(inp: VmecInput, max_mode: int, *, vary_major_radius: bool = False) → ndarray

Flat boundary-dof vector (see _dof_modes()).

Inverse of unpack_boundary(); RBC(0,0) is excluded by default (fixed major radius) and appended when vary_major_radius=True. 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, *, vary_major_radius: bool = False) → 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.residuals_from_tuples(state: SpectralState, runtime: SolverRuntime, objective_terms: Sequence[tuple[Callable, float, Any]], *, weight_semantics: str = 'cost') → Array

Stack traceable (function, target, weight) objective rows.

With the default weight_semantics="cost", each tuple contributes sqrt(weight) * (function(state, runtime) - target) so the usual scalar objective is simply 0.5 * residuals @ residuals. This small public building block is useful when a user owns the equilibrium map, such as a differentiable free-boundary solve, and wants to pass the resulting value and gradient to SciPy, JAXopt, Optax, or a custom optimizer directly.

vmex.core.optimize.resample_current_profile(inp: VmecInput, n_spline: int, *, kind: str = 'cubic_spline_ip') → VmecInput

Represent the current shape on n_spline uniform spline knots.

The existing enclosed-current profile is differentiated and sampled at the new knots, so continuation preserves I(s) before the new knot values are optimized. CURTOR remains the independent amplitude. kind may be cubic_spline_ip or akima_spline_ip.

vmex.core.optimize.least_squares(objective_terms: Sequence[tuple[Callable, float, Any]], inp: VmecInput, *, max_mode: int | Sequence[int] = 1, vary_major_radius: bool = False, x0: ndarray | None = None, current_dofs: int | None = None, jac: str | None = None, jac_chunk_size: int | str | None = 'auto', jac_solver: str = 'auto', adjoint_tol: float = 1e-06, jacobian_adjoint_tol: float = 0.0001, jacobian_adjoint_maxiter: int = 10, adjoint_maxiter: int = 300, max_fsq_ratio: float = 1000000.0, refine_tol: float = 1e-10, forward_ftol: float | None = None, forward_max_iterations: int | None = None, hot_restart: bool = True, warm_start: str | None = 'perturbation', use_ess: bool = False, ess_alpha: float = 1.2, device: Any = 'auto', 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). 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. Staged campaigns with different objectives are just successive calls, each seeded with the previous call’s result.input.

current_dofs = k additionally frees the current profile: the first k AC coefficients, or the first k AC_AUX_F values for a spline profile, plus CURTOR. For an n-knot spline, use current_dofs=n-1: the remaining fixed ordinate removes the profile’s overall-scale null direction because CURTOR already sets that scale. The values are scaled by their frozen seed magnitude so the trust region sees O(1) numbers. Requires ncurr = 1; resample_current_profile() changes spline resolution between continuation stages without changing the represented I(s). Both gradient modes support it (finite differences re-solve per current dof; jac="implicit" adds k + 1 one-hot tangent rows through ImplicitParams.ac or ac_aux_f and curtor). VMEC normalizes the AC profile by its own edge integral (only the shape of I' matters; CURTOR sets the amplitude). This is the dof set of 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. Trial solves are cheap by construction: runtimes with the same Resolution are structural pytrees, so one XLA executable is reused across all boundary trials (only the first solve of a stage compiles). device is forwarded to the solver ("auto" applies vmex.core.device’s policy; None follows JAX placement). use_ess enables Exponential Spectral Scaling of the trust region (_ess_scale(), ess_alpha).

jac=None (default) uses scipy "2-point" finite differences. jac="implicit" computes the exact residual Jacobian by forward implicit differentiation (module docstring): one hot-restarted forward solve per trial boundary and one preconditioned GMRES per boundary dof instead of one full equilibrium solve per dof. Every term must be 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 — use the traceable mercier_stability_residual() / l_grad_b_state() instead). Symmetric and lasym boundaries are both supported. The knobs below are inert for jac=None.

jac_chunk_size chunks the per-dof Jacobian columns via solvax.chunk_map(): "auto" (default) caps SOLVAX’s device-aware width by a conservative square-root policy, so an accelerator memory report cannot expand the full probe batch; an int fixes that many dofs at a time; None forces one wide batch. Automatic widths prefer a nearby divisor to avoid compiling a separate tail. Column blocks are mathematically independent, so the assembled Jacobian is identical across chunk sizes to float64 round-off.

jac_solver selects the implicit-Jacobian direction. "auto" (default) uses one matrix-free reverse solve for a scalar residual and the "block" path otherwise. "reverse" requests one reverse solve per residual row. "block" 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 chunked three-surface VJPs capture it completely at a cost independent of the dof count — then backsolves every dof right-hand side (solvax.block_thomas_factor() / solvax.block_thomas_solve()) and certifies each column with a warm-started GMRES pass against the preconditioned system (same adjoint_tol and configured iteration budget; columns already at tolerance cost one matvec). "gmres" is the per-dof-column fallback if the block path misbehaves on an exotic configuration; both produce the same Jacobian to solver tolerance. An uncertified automatic response is recomputed by reverse adjoint. Forced host methods raise AdjointSolveError; transformable JAX methods fall back to reverse so no compiled public path returns an uncertified matrix.

adjoint_tol and adjoint_maxiter control the certified linear response solves. All public optimization paths allow 300 restarts by default so QI derivatives cannot fail merely because of the former hard-coded 30-restart cap. Easier problems normally stop far earlier.

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 (jac="implicit" only) refines that seed. "perturbation" (default) seeds each trial 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 are exactly the columns the implicit Jacobian already solves, so the linearization is stashed at each jac(x_ref) call for free. "state" is the plain hot restart; None disables warm starting. All three converge to the same fixed points — only the inner iteration count changes — and a missing or mismatched seed falls back through the perturbation -> state -> cold ladder. hot_restart=False forces warm_start=None.

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 (cumulative solve, refinement, Jacobian and adjoint counters of the stage’s configuration, see implicit._SOLVE_STATS).

vmex.core.optimize.minimize(objective_terms: Sequence[tuple[Callable, float, Any]], inp: VmecInput, *, max_mode: int | Sequence[int] = 1, vary_major_radius: bool = False, x0: ndarray | None = None, current_dofs: int | None = None, hot_restart: bool = True, device: Any = 'auto', solve_kwargs: dict | None = None, verbose: int = 0, method: str = 'L-BFGS-B', adjoint_tol: float = 1e-06, adjoint_maxiter: int = 300, max_fsq_ratio: float = 1000000.0, refine_tol: float = 1e-10, forward_ftol: float | None = None, forward_max_iterations: int | None = None, **scipy_kwargs)

Minimize the scalarized residual norm with one adjoint per gradient.

The objective is exactly 0.5 * sum(rows**2), with rows defined by least_squares(). Unlike Gauss–Newton least squares, a reverse gradient of this scalar needs one matrix-free implicit adjoint and never forms the vector residual Jacobian or its dense radial block factors. This is the bounded-storage path for profile objectives such as DMerc, jdotb, and Glasser D_R. It changes the optimization algorithm, not the objective or its unconstrained minimizers, and is therefore opt-in; least_squares() retains all existing defaults.

method and remaining keywords are passed to scipy.optimize.minimize() (default "L-BFGS-B"; use bounds= and options={"maxiter": ...} in the usual scipy form). All objective terms must support jac="implicit" as documented by least_squares(). Plain state hot restarts are used because the first-order perturbation warm start requires the forward state-response columns that this lower-storage path deliberately avoids.

adjoint_tol / adjoint_maxiter are exposed explicitly. Their defaults are certified on the QI, QS, L_grad_B, DMerc and D_R objective lanes; an unconverged adjoint is never returned as a gradient.

The optimizer-neutral callables behind vj.VmecProblem.from_tuples in the README: value, residual, and derivative functions with the contracts SciPy, JAXopt, Optax, and user code consume, and no optimization algorithm of their own.

Optimizer-neutral objective and derivative callables.

The classes in this module contain no optimization algorithms. They expose the small function contracts consumed by SciPy, JAXopt, Optax, and user code. VMEC-specific construction is imported lazily so this module remains usable in lightweight tests and does not introduce an import cycle with vmex.core.optimize.

class vmex.core.problem.Evaluation(x: ndarray, value: float | None = None, gradient: ndarray | None = None, residual: ndarray | None = None, jacobian: ndarray | None = None, status: str = 'success', message: str = '', diagnostics: Mapping[str, ~typing.Any]=<factory>)

Values and diagnostics produced at one decision vector.

Fields that were not requested or are unavailable are None. status is a short machine-readable value such as "success" or "failed_solve"; message is intended for a human. Optimizers use the ordinary callable methods and do not need to understand this object.

Returned by FunctionProblem.evaluate() and the two compile_* helpers. It is a report, not a cache: nothing here is consulted by a later evaluation.

x

The decision vector the values were produced at, as an owned float copy with the shape of problem.x0.

Type:

numpy.ndarray

value

Scalar objective at x. For a residual-only problem this is the least-squares cost 0.5 * r @ r, matching SciPy’s OptimizeResult.cost. None when the problem exposes neither a scalar objective nor residuals.

Type:

float | None

gradient

Gradient of value with respect to x, reshaped to x’s shape. For a residual problem it is J.T @ r. None when derivatives=False was requested or no gradient lane exists.

Type:

numpy.ndarray | None

residual

The flattened residual vector r(x), or None for a scalar-only problem.

Type:

numpy.ndarray | None

jacobian

The residual Jacobian dr_i/dx_j with shape (residual.size, x.size), or None when derivatives were not requested or the problem provides no Jacobian.

Type:

numpy.ndarray | None

status

"success", or — from a VMEC-backed problem — "failed_solve" when the equilibrium solve at x raised, and "under_converged" when it returned but its force residual exceeds the threshold below which implicit derivatives are certified.

Type:

str

message

Human-readable explanation, empty on success; the solver exception’s text for "failed_solve".

Type:

str

diagnostics

Extra per-evaluation values. Empty for a plain FunctionProblem. A VmecProblem adds the cumulative failed_trials and derivative_fallbacks counters, a solve_stats mapping when the implicit lane recorded one, and — when the equilibrium at x could be materialised — the summed force residual fsq, its fsq_ratio to the solve tolerance, the configured max_fsq_ratio, and the boolean derivative_certified. A failed solve also carries exception_type.

Type:

Mapping[str, Any]

property success: bool

Whether the evaluation completed without a recoverable failure.

class vmex.core.problem.FunctionProblem(x0: Any, *, fun: Callable[[ndarray], Any] | None = None, grad: Callable[[ndarray], Any] | None = None, value_and_grad: Callable[[ndarray], Any] | None = None, residual: Callable[[ndarray], Any] | None = None, residual_jac: Callable[[ndarray], Any] | None = None, residual_and_jac: Callable[[ndarray], Any] | None = None, jax_fun: Callable[[Any], Any] | None = None, jax_value_and_grad: Callable[[Any], tuple[Any, Any]] | None = None, jax_residual: Callable[[Any], Any] | None = None, jax_residual_jac: Callable[[Any], Any] | None = None, names: Sequence[str] | None = None, bounds: Any = None, scales: Any | None = None, metadata: Mapping[str, Any] | None = None, evaluation_progress: bool = False, report_interval: float = 10.0)

A decision vector plus optimizer-compatible objective callables.

Parameters are explicit and immutable from the caller’s perspective. Supplying combined value/gradient or residual/Jacobian functions enables a one-entry exact-key cache, so the common SciPy call sequence does not repeat expensive work. The cache is protected by a lock; JAX-native callables do not use host state and remain suitable for tracing.

This class deliberately does not provide solve(method=...). Pass its methods directly to the optimizer of choice.

At least one of fun, value_and_grad, residual, or residual_and_jac is required; every other callable is optional and the matching method raises AttributeError when its lane is absent. The host callables receive one contiguous float NumPy array and are free to be opaque; the jax_* callables receive whatever the caller traces and must stay traceable.

Two independent one-entry caches, each keyed on the exact bytes of x (shape, dtype, and contents — no tolerance), avoid repeating work across the split calls an optimizer makes at one iterate. The scalar cache is filled by value_and_grad() and so covers value_and_grad alone, fun together with grad, or residual_and_jac. The least-squares cache is filled by residual_and_jac() and so covers residual_and_jac, or residual together with residual_jac. Whether the cache actually pays depends on which callables were supplied: with residual_and_jac, SciPy’s separate fun(x) then jac(x) calls both route through it and the second is free, whereas separately supplied residual and residual_jac are each invoked directly and share nothing. Cached arrays are copied out, so a caller may mutate what it receives. The caches are guarded by a re-entrant lock; the jax_* lane touches no host state and stays safe to trace.

Parameters:

  • x0 – Initial decision vector. Copied to a float array; its size fixes the number of degrees of freedom and the expected Jacobian column count, and its shape is the shape gradients are reshaped to.

  • fun – x -> float, the scalar objective. Called directly by fun() without touching the cache.

  • grad – x -> array, the objective gradient. Used only in combination with fun; on its own it does not enable grad().

  • value_and_grad – x -> (float, array). The preferred scalar lane: it is one call for both quantities and it fills the scalar cache. The method that serves it is also reachable under SciPy’s fun_and_grad name.

  • residual – x -> array, the least-squares residual vector r(x). It is flattened, and defines the scalar objective 0.5 * r @ r when the problem has no scalar lane of its own (no fun, grad, or value_and_grad).

  • residual_jac – x -> array, the Jacobian dr_i/dx_j. It must have one column per decision variable; anything else raises ValueError.

  • residual_and_jac – x -> (array, array), the preferred least-squares lane: one call for both, filling the cache that residual() and residual_jac() then read. The Jacobian shape is checked exactly against (r.size, x0.size).

  • jax_fun – Traceable counterparts of the four callables above, returned unwrapped by the matching jax_* methods. They are never cached and never see host state, so they remain usable inside jax.jit() and jax.grad(). jax_fun falls back to the first element of jax_value_and_grad when it is not supplied.

  • jax_value_and_grad – Traceable counterparts of the four callables above, returned unwrapped by the matching jax_* methods. They are never cached and never see host state, so they remain usable inside jax.jit() and jax.grad(). jax_fun falls back to the first element of jax_value_and_grad when it is not supplied.

  • jax_residual – Traceable counterparts of the four callables above, returned unwrapped by the matching jax_* methods. They are never cached and never see host state, so they remain usable inside jax.jit() and jax.grad(). jax_fun falls back to the first element of jax_value_and_grad when it is not supplied.

  • jax_residual_jac – Traceable counterparts of the four callables above, returned unwrapped by the matching jax_* methods. They are never cached and never see host state, so they remain usable inside jax.jit() and jax.grad(). jax_fun falls back to the first element of jax_value_and_grad when it is not supplied.

  • names – One name per decision variable, in order, surfaced as dof_names for labelling output. The default is x[0], x[1], ...; a length other than x0.size raises ValueError.

  • bounds – Box constraints stored verbatim for the optimizer to consume — a SciPy Bounds, or a (lower, upper) pair. This class neither interprets nor enforces them.

  • scales – Positive finite per-variable scale factors with the shape of x0, defaulting to ones. Like bounds they are carried, not applied: pass them to the optimizer (SciPy’s x_scale). A non-finite or non-positive entry raises ValueError.

  • metadata – Free-form mapping copied onto the instance. VMEX-built problems use it to carry the named residual slices, the solver configuration, the mutable solve counters, and the traceable state accessors that VmecProblem reads.

  • evaluation_progress – Print an elapsed-time heartbeat around long evaluations. It stays silent until a call outlives the first interval, so fast calls print nothing. It wraps the standalone residual() and residual_jac() calls only, which is where a production deck spends minutes; the combined and scalar lanes are unaffected.

  • report_interval – Seconds between heartbeat lines. Must be positive.

property dof_names: tuple[str, ...]

Ordered names corresponding one-to-one with entries of a decision vector.

classmethod from_functions(x0: Any, **kwargs: Any) → FunctionProblem

Build a problem from user-supplied x-level callables.

value_and_grad(x: Any) → tuple[float, ndarray]

Return scalar value and gradient for SciPy jac=True.

fun_and_grad(x: Any) → tuple[float, ndarray]

Return scalar value and gradient for SciPy jac=True.

fun(x: Any) → float

Return the scalar objective value.

grad(x: Any) → ndarray

Return the scalar objective gradient.

residual_and_jac(x: Any) → tuple[ndarray, ndarray]

Return the residual vector and its Jacobian.

residual(x: Any) → ndarray

Return the residual vector.

residual_jac(x: Any) → ndarray

Return the residual Jacobian.

J(x: Any) → float

SIMSOPT-style alias for fun().

dJ(x: Any) → ndarray

SIMSOPT-style alias for grad().

jax_fun(x: Any) → Any

Return the traceable scalar objective.

jax_value_and_grad(x: Any) → tuple[Any, Any]

Return the traceable scalar value and gradient.

jax_residual(x: Any) → Any

Return the traceable residual vector.

jax_residual_jac(x: Any) → Any

Return the traceable residual Jacobian.

evaluate(x: Any, *, derivatives: bool = True) → Evaluation

Evaluate available scalar and residual quantities at x.

compile_residual_and_jacobian(x: Any | None = None, *, progress: bool = True, report_interval: float = 10.0, stream: Any = None) → Evaluation

Compile and cache the least-squares residual and Jacobian.

This call is optional: an optimizer compiles on its first evaluation if it is omitted. Calling it explicitly provides elapsed-time output during a potentially long first JAX compilation. Later calls at the same x use the normal one-entry problem cache.

compile_value_and_gradient(x: Any | None = None, *, progress: bool = True, report_interval: float = 10.0, stream: Any = None) → Evaluation

Compile and cache the scalar value and gradient.

This optional call makes the first JAX compilation visible before BFGS, L-BFGS-B, Adam, or another gradient optimizer starts.

class vmex.core.problem.VmecProblem(*args: Any, input_from_x: Callable[[Any], Any], x_from_input: Callable[[Any], Any], equilibrium_from_x: Callable[[...], Any] | None = None, boundary_from_x: Callable[[Any], Any] | None = None, **kwargs: Any)

A FunctionProblem backed by a VMEX equilibrium solve.

Adds the maps between the optimizer’s decision vector and VMEC objects: the input deck, the converged equilibrium, and the boundary coefficient arrays. It keeps the same optimizer contract as its base class, so the same methods go to SciPy, JAXopt, or Optax unchanged, and it enriches evaluate() with the solve and adjoint status of the underlying equilibrium.

Build one with from_tuples(), from_loss(), or from_input() rather than calling this constructor: they route through vmex.core.optimize.make_problem, which is what assembles the four callables below along with the objective, the derivative lane, the degree-of-freedom names, and the metadata.

Parameters:

  • *args – Forwarded positionally to FunctionProblem; in practice the decision vector x0.

  • **kwargs – Forwarded to FunctionProblem: the objective callables, names, bounds, scales, and metadata.

  • input_from_x – Required. x -> VmecInput: a new input deck carrying the boundary coefficients — and the current degrees of freedom, when the problem parameterizes them — of this decision vector. Nothing is solved.

  • x_from_input – Required. The inverse, VmecInput -> array: the decision vector that reproduces a given deck, which is the normal starting point of a continuation stage. x_from_input() rejects a result whose shape differs from x0.

  • equilibrium_from_x – Optional x -> Equilibrium, the converged equilibrium at x. Implicit problems return the accepted state the objective already computed rather than cold-solving the boundary again, which matters for strongly shaped boundaries whose cold axis guess can produce a sign-changing initial Jacobian. None makes equilibrium_from_x() raise AttributeError. A callable that accepts a newton_iterations keyword receives it only when the caller asks for something other than the default 10, so a closure without that keyword still works.

  • boundary_from_x – Optional x -> tuple of arrays, the traceable boundary coefficients: (rbc, zbs) for a stellarator-symmetric input and (rbc, zbs, rbs, zbc) when lasym. Each is a full dense INDATA-layout array of shape (2 * ntor + 1, mpol) indexed [n + ntor, m] in metres — not the trimmed decision vector — and is a JAX array, so it composes with coil or surface objectives under jax.grad(). None makes boundary_from_x() raise AttributeError.

classmethod from_tuples(inp: Any, objective_terms: Sequence[tuple[Callable[[...], Any], Any, float]], **kwargs: Any) → VmecProblem

Build a VMEC least-squares problem from weighted objective tuples.

The README entry point. Each tuple is (function, target, weight) and contributes one or more rows to a single residual vector; the rows of all terms are concatenated in the order given and the scalar cost is 0.5 * r @ r. Named row ranges are recorded in the problem metadata, which is what lets OptimizationMonitor report per-term costs without re-solving anything.

Parameters:

  • inp – The starting VmecInput. Its boundary supplies the initial decision vector and its resolution and profiles are held fixed apart from the parameterized degrees of freedom.

  • objective_terms –

    The (function, target, weight) triples.

    function is normally a traceable function(state, runtime) -> scalar or vector, evaluated on the converged equilibrium — the two arguments are also spelled (equilibrium_state, solver_context). An objective object exposing a residuals_state method may be passed instead (whole instance or bound method), in which case its full pointwise residual vector becomes this term’s rows. Under derivative_method="finite_difference" a one-argument host callable taking the whole Equilibrium is accepted too; the implicit lane rejects it, since it cannot be traced.

    target is the value the term is driven toward, coerced with float(), so it must be scalar even when function returns a vector — the same target is then subtracted from every row.

    weight is a non-negative scalar, or a one-dimensional array with one entry per residual row of that term. Under the default weight_semantics="cost" it multiplies the squared cost, so the row is sqrt(weight) * (function - target); with weight_semantics="residual" the row is weight * (function - target) and a negative entry is then allowed.

  • **kwargs – Passed through to vmex.core.optimize.make_problem: which boundary modes vary (max_mode, vary_major_radius, current_dofs), the derivative lane (derivative_method, implicit_jacobian_method, jacobian_batch_size), the forward solve controls, the variable scaling (use_ess, ess_alpha, bounds), and weight_semantics.

Returns:

  • A VmecProblem whose residual/residual_jac pair,

  • x0, and scales are ready for

  • scipy.optimize.least_squares(). A non-finite or empty

  • residual at the initial point raises FloatingPointError

  • rather than starting an optimization that cannot recover.

classmethod from_loss(inp: Any, loss: Callable[[...], Any], **kwargs: Any) → VmecProblem

Build a VMEC scalar problem from a traceable state/runtime loss.

Parameters:

  • inp – The starting VmecInput, as for from_tuples().

  • loss – loss(state, runtime) -> scalar, evaluated on the converged equilibrium and already carrying its own weights. It must return a single value: a vector-valued objective belongs in from_tuples(), or must be reduced here explicitly. Unlike an objective_terms entry it is used exactly as written — an object’s residuals_state is never substituted for it.

  • **kwargs – As for from_tuples().

Returns:

  • A VmecProblem exposing only the scalar lane —

  • fun(), grad(), and

  • their traceable counterparts — for a gradient optimizer such as

  • BFGS, L-BFGS-B, or Adam. It provides no residual or Jacobian.

classmethod from_input(inp: Any, **kwargs: Any) → VmecProblem

Parameterize an input for field VJPs without defining an objective.

Builds the same machinery as from_loss() around an identically zero loss, so there is nothing to minimize. Use it when what you want is the parameterization itself: the decision vector and its names, input_from_x() and boundary_from_x(), and the differentiable interior_field() and exterior_field() with exact VJPs in these degrees of freedom. inp and **kwargs are as for from_tuples().

input_from_x(x: Any) → Any

Return a new VmecInput containing decision vector x.

x_from_input(inp: Any) → ndarray

Return this problem’s decision vector for inp.

This is the inverse of input_from_x() for the boundary and any optional current degrees of freedom selected when the problem was constructed. It is the normal continuation-stage starting vector.

equilibrium_from_x(x: Any, *, newton_iterations: int = 10) → Any

Return the converged equilibrium evaluated at x.

Implicit problems reuse the accepted optimizer state instead of cold-solving the optimized boundary again. This matters for strongly shaped boundaries whose cold magnetic-axis guess may have a sign-changing initial Jacobian.

boundary_from_x(x: Any) → Any

Return traceable boundary coefficient arrays for decision vector x.

jax_objective_from_state(x: Any, extra_costs: Callable[[Any, Any], Any], *, n_extra_terms: int) → tuple[Any, tuple[Any, Any]]

Combine the VMEX least-squares cost with state-dependent costs.

extra_costs(state, runtime) returns one already-weighted scalar cost per added objective term. The auxiliary result contains the VMEX residual rows and those added costs, ready to pass as auxiliary data to jax.value_and_grad(). Failed equilibrium trials receive the same smooth finite rejection cost as the base problem, so driver scripts do not need their own accepted/rejected branches.

jax_extra_costs_from_state(x: Any, extra_costs: Callable[[Any, Any], Any], *, n_extra_terms: int) → tuple[Any, Any]

Evaluate additive state-dependent costs only at valid VMEC trials.

This is the split-compilation counterpart of jax_objective_from_state(). It returns zero extra cost at a rejected trial, leaving the base problem to supply its certified rejection wall. Splitting a large virtual-casing or coil graph from the VMEC objective substantially lowers peak XLA compilation memory.

jax_quantity_from_state(x: Any, quantity: Callable[[Any, Any], Any]) → tuple[Any, Any]

Evaluate a differentiable floating-point quantity and solve status.

quantity(state, runtime) may return any fixed-shape JAX array. Rejected equilibrium trials return an array of NaNs with that shape, so an invalid state cannot look like a usable diagnostic. The scalar status is zero only for an accepted equilibrium.

exterior_field(x: Any, *, external_field: Any | None = None, external_parameters: Any | None = None, external_field_from_parameters: Callable[[Any], Any] | None = None, external_dof_names: tuple[str, ...] = (), nphi: int = 32, ntheta: int = 32, digits: int = 6, levels: tuple[tuple[int, int], ...] | None = None, chunk_size: int | str = 'auto', target_chunk_size: int | str = 'auto') → Any

Return the exterior field and exact VJPs in this problem’s DOFs.

Query points must lie outside the last closed flux surface and away from coil filaments. The returned field follows the stored-point API: field.set_points(xyz); field.B(); field.B_vjp(cotangent). Set the source and target chunk sizes only to cap virtual-casing memory; "auto" is the tuned default.

interior_field(x: Any, *, newton_iterations: int = 10) → Any

Return the interior field and exact VJPs in this problem’s DOFs.

surface_field_values(x: Any, quantity: str, *, external_field: Any | None = None, nphi: int = 32, ntheta: int = 32, digits: int = 4, precision: Any | None = None) → Any

Return |B| or B.n/B on a trial boundary for plotting.

B.n/B is evaluated on the exterior side using the supplied coil or MGRID field plus the plasma-current virtual-casing field. This helper keeps optional movie coloring out of optimization driver code; it is not used by the objective or optimizer.

evaluate(x: Any, *, derivatives: bool = True) → Evaluation

Evaluate and attach VMEC solve/adjoint status diagnostics.

Accepted-iteration reporting independent of optimization algorithms.

class vmex.core.monitoring.EquilibriumReporter(*quantities: tuple[str, ~typing.Callable[[...], ~typing.Any], str], stream: ~typing.TextIO | None | object = <object object>, separator: str = ', ')

Print a compact set of scalar diagnostics for an equilibrium.

Each quantity is (label, callable, format_spec). Callables may use either the function(equilibrium) or function(state, runtime) convention used by VMEX objectives. Calling the reporter prints one line and returns the values by label, so scripts can also reuse a final metric.

Parameters:

  • *quantities – One (label, function, format_spec) triple per reported column; at least one is required and the labels must be unique. label names the column and keys the returned mapping. function is dispatched on its signature: a callable whose second positional parameter has no default is called as function(state, runtime) (the VMEX objective convention), every other callable as function(equilibrium). It must return exactly one scalar; any other size raises ValueError. format_spec is a format() specification applied to that float, for example ".6e".

  • stream – Where the report line is written. The default is sys.stdout; pass None to compute and return the values without printing.

  • separator – Text placed between the label = value fields of the printed line.

class vmex.core.monitoring.OptimizationMonitor(problem: FunctionProblem | None = None, *, stream: TextIO | None | object = <object object>, print_every: int = 1, trace: bool = True)

Record and optionally print optimizer iterations and trials.

Pass the instance as a SciPy callback. SciPy invokes callbacks after an iteration, unlike objective functions which are also called for rejected line-search or trust-region trials. JAXopt, Optax, and custom loops can call record() with values they already computed. trace also prints one trial line per objective evaluation, so the rejected line-search trials between two accepted iterations – a full equilibrium solve each – are visible while the run is in progress.

The monitor never chooses steps or changes an optimizer. If problem is supplied, VMEX solve/failure counters are read without evaluating the objective again.

Parameters:

  • problem – Optional problem the optimizer is running on. It is used only to read metadata: the named residual slices that split cost into per-term costs, the cumulative equilibrium-solve and failed-trial counters, and — as a last resort, when a callback carries neither cost nor fun — one fun() call at the accepted iterate. With None those record fields stay None or empty and the monitor still records everything the callback provides.

  • stream – Where the per-iteration table is written. The default is sys.stdout; pass None to record silently and read records, history, or save() afterwards.

  • print_every – Print one row every print_every records (the first record is always printed, together with the header). Must be at least 1; recording is unaffected.

wrap_value_and_grad(function: Callable | Sequence[Callable], term_names: tuple[str, ...] | None = None, *, residual_slices: tuple[tuple[str, int, int], ...] = ()) → Callable

Adapt a JAX has_aux value/gradient pair for SciPy.

Each function(x) must return ((cost, terms), gradient). Pass a sequence to compile large additive physics components separately; their costs and gradients are summed without changing the optimizer contract. terms may be a mapping of labels to weighted scalar costs, or one compact vector paired with term_names. For a large residual graph, pass residual_slices and return (residual, extra_costs...); costs are reduced on the host to keep the compiled output small. The first evaluation is recorded as iteration zero; later evaluations are cached.

cache_evaluation(x: Any, cost: Any, gradient: Any, terms: Mapping[str, Any] | None = None) → tuple[float, ndarray]

Cache an already-computed objective pair for a SciPy callback.

This method does no differentiation and does not alter the objective. Driver scripts can show their explicit jax.value_and_grad calls, sum independently compiled physics components themselves, and use this one host conversion to avoid recomputing per-term costs when SciPy later reports an accepted iterate.

record(x: Any, *, cost: float, optimality: float | None = None, iteration: int | None = None, equilibrium_solves: int | None = None, rejected_trials: int | None = None, terms: Mapping[str, Any] | None = None, counters: Mapping[str, Any] | None = None) → OptimizationRecord

Append one already-computed accepted iterate and return its record.

Parameters:

  • x – The accepted decision vector; a float copy is appended to x_history.

  • cost – Total scalar objective at x. Nothing is recomputed.

  • optimality – First-order optimality measure, or None when unknown.

  • iteration – Iteration index; defaults to the number of records already held, and is advanced past the previous record when an optimizer restarts its own counter at a continuation stage.

  • equilibrium_solves – Cumulative solve and failed-trial counts. None (the default) reads them from the monitor’s problem, leaving the field None when no problem was supplied.

  • rejected_trials – Cumulative solve and failed-trial counts. None (the default) reads them from the monitor’s problem, leaving the field None when no problem was supplied.

  • terms – Per-term weighted costs by label. None splits the problem’s own residual over its named term slices instead, which may evaluate the residual once at x.

  • counters – Effort counters by name. None reads them from the monitor’s problem, empty when it carries none.

Return type:

The appended OptimizationRecord.

property history: dict[str, ndarray]

Return total and per-term cost histories as one mapping.

property x_history: tuple[ndarray, ...]

Copies of the accepted decision vectors, including iteration zero.

save(path: str | Path) → Path

Save the recorded iteration and cost columns as CSV.

plot(path: str | Path, *, title: str = 'Optimization objective terms') → Path

Write a compact log-scale total and per-term cost history plot.

movie(path: str | Path, object_factory: Callable[[ndarray], Any], **kwargs: Any) → Path

Animate accepted surface/coil iterates with one geometry callback.

movie_surface_coils(path: str | Path, object_factory: Callable[[ndarray], Any], *, x0: Any, scales: Any, surface_color: str | Callable | None = None, plasma_problem: Any | None = None, external_field: Callable[[Any], Any] | None = None, nphi: int = 32, ntheta: int = 32, digits: int = 4, precision: Any | None = None, **kwargs: Any) → Path

Animate normalized surface/coil iterates with optional field color.

object_factory receives physical variables x0 + scales*u. surface_color may be None, "absB", "B.n/B", or a callable (u, objects) -> values. For the two named field colors, provide plasma_problem; B.n/B additionally needs external_field(objects). These plotting helpers never enter the objective or gradient graph.

class vmex.core.monitoring.OptimizationRecord(iteration: int, cost: float, reduction: float | None, optimality: float | None, equilibrium_solves: int | None, rejected_trials: int | None, terms: Mapping[str, float]=<factory>, counters: Mapping[str, float | None]=<factory>)

One optimizer callback, normally one accepted iteration.

Produced by OptimizationMonitor.record() and stored in OptimizationMonitor.records. Every field is a plain host value; a field is None when the optimizer callback did not supply it and the monitor could not derive it, never zero-as-unknown.

iteration

Iteration index of this accepted iterate. Taken from the SciPy result’s nit when present, otherwise the number of records already held. Optimizers restart their counter at every continuation stage, so a repeated or decreasing value is bumped to one past the previous record; the sequence in one monitor is always strictly increasing.

Type:

int

cost

Total scalar objective at this iterate. For a least-squares problem this is 0.5 * r @ r over the full residual vector, so it matches SciPy’s OptimizeResult.cost; for a scalar objective it is the objective value itself. Units are those of the weighted objective (dimensionless for the usual normalised VMEX terms).

Type:

float

reduction

previous cost - this cost, positive when the step improved the objective. None for the first record, which has no predecessor.

Type:

float | None

optimality

First-order optimality measure of the gradient. SciPy’s own optimality is used when the callback provides it; otherwise it is the infinity norm of a callback-supplied jac, or the Euclidean norm of the gradient cached by OptimizationMonitor.cache_evaluation(). None when no gradient information reached the monitor.

Type:

float | None

equilibrium_solves

Cumulative number of forward VMEC equilibrium solves performed for this problem’s implicit configuration, read from the solver’s own counters without re-evaluating the objective. None when the monitor was built without a problem or the problem carries no implicit configuration (for example a finite-difference or wout-only problem).

Type:

int | None

rejected_trials

Cumulative number of optimizer trial points whose equilibrium solve failed and was replaced by the smooth rejection cost. None under the same conditions as equilibrium_solves.

Type:

int | None

terms

Per-term weighted costs by label, 0.5 * r_k @ r_k over each named residual slice of the problem, so the values sum to cost for a pure least-squares problem. Empty when the problem exposes no term slices and the caller supplied none.

Type:

Mapping[str, float]

counters

Cumulative effort counters of the problem’s implicit configuration at this record: solves and descent iterations, refinement calls, steps and Krylov iterations, Jacobian calls, columns and certifier Krylov iterations, adjoint calls and Krylov iterations, and the host seconds of each part. An entry is None when the work ran inside a compiled program where it cannot be observed. Empty under the same conditions as equilibrium_solves.

Type:

Mapping[str, float | None]

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 (10-core host, 8 balanced phiedge solves): 1.79x with 2 workers, 3.29x with 8. 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 Why threading, not pmap 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.available_cpus() → int

Return the number of CPUs available, respecting scheduler limits.

vmex.core.parallel.default_workers(n_items: int, workers: int | None = None) → int

Resolve the worker count for an n_items ensemble.

None uses the CPUs available to this process, including scheduler affinity limits. An explicit value is clamped to [1, n_items].

vmex.core.parallel.evaluate_problems(problems: Sequence[Any], xs: Sequence[Any] | None = None, *, derivatives: bool = True, workers: int | None = None, return_exceptions: bool = False) → list[Any]

Evaluate independent problem objects concurrently in input order.

Use one problem per ensemble or multistart member so each equilibrium cache remains local. xs defaults to each problem’s x0.

vmex.core.parallel.finite_difference_gradient(fun: Callable[[Any], Any], x: Any, **kwargs: Any) → Any

Parallel finite-difference gradient of a scalar host function.

vmex.core.parallel.finite_difference_jacobian(fun: Callable[[Any], Any], x: Any, *, method: str = '3-point', rel_step: float | None = None, workers: int | None = None) → Any

Differentiate an opaque host function with independent parallel probes.

method="3-point" (default) is central and second-order accurate; "2-point" is forward and first-order accurate. workers=None uses default_workers(), while workers=1 is deterministic serial execution. Every probe receives its own copy of x and fun must not mutate shared state.

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 Why threading, not pmap.

Outputs

VMEC2000-compatible wout_*.nc schema, writer and reader.

This module implements the full variable set written by VMEC2000’s wrout.f (Appendix A). Core fixed-boundary fields use the reference netCDF names, dimensions, dtypes and unit conventions. Free-boundary results also populate the symmetric or LASYM NESTOR potential and surface fields when their public VacuumOutput is supplied:

  • 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. VMEX-written files also carry the versioned vmex_trapped_fraction radial profile; older WOUT files remain readable and are rewritten without the extension.

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, vmex_diagnostics_schema: int = 0, vmex_trapped_fraction: 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, vacuum_output: VacuumOutput | None = None) → WoutData

Build a complete WoutData from a solved equilibrium 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). Pass the public result.vacuum as vacuum_output to populate NESTOR potsin/xmpot/xnpot and *_sur tables. LASYM runs additionally populate potcos and the four sine *_sur partners.

Optional neoclassical diagnostics

Effective-ripple diagnostics through the optional NEO_JAX backend.

vmex.core.neoclassical.diagnostic_neo_config()

Return the bounded NEO resolution used by VMEX summary figures.

On the bundled NFP=2 QA and NFP=4 QI finite-beta decks this stays within 4% of NEO_JAX’s default NeoConfig. A 16 x 16 spline grid aliases the summary Boozer spectrum (mboz=16) and fewer than 20 steps per field period or 200 field periods leave 7-30% errors. Pass a neo_jax.NeoConfig to epsilon_effective_from_wout() for publication calculations.

vmex.core.neoclassical.epsilon_effective_from_boozer(booz: Any, *, config=None)

Return (s, epsilon_eff**(3/2)) from Boozer-coordinate data.

booz may be a NEO_JAX BoozerData or a booz_xform-style mapping. The returned NEO quantity is the effective-ripple transport measure \(\epsilon_\mathrm{eff}^{3/2}\), conventionally named epstot. JAX arrays remain differentiable when the mapping and selected NEO path are JAX-native.

vmex.core.neoclassical.epsilon_effective_from_wout(wout, *, surfaces: Sequence[float] = (0.2, 0.4, 0.6, 0.8, 0.95), mboz: int = 16, nboz: int = 12, config=None, clear_jax_caches: bool = False)

Compute epsilon_eff**(3/2) directly from an in-memory VMEX wout.

VMEX performs the Boozer transform in memory and passes its arrays to NEO_JAX; no boozmn file is required. NEO_JAX currently represents the stellarator-symmetric cosine/sine convention, so LASYM wouts are rejected rather than silently dropping asymmetric harmonics. The default is library-safe: a diagnostic call never clears the process-wide JAX executable caches behind its caller’s back. Pass clear_jax_caches=True to release completed executables before compiling NEO when peak memory matters more than warm executables — the CLI does this itself after all requested diagnostics instead (vmex.core.cli).

Optional alpha-particle tracing

The vmex-to-ESSOS field handoff, and alpha tracing on top of it.

  • essos_vmec_field() — hand a solved equilibrium (or a wout file) to ESSOS as an essos.fields.Vmec, ready for ESSOS tracing, surfaces and field queries. The seam runs one way, ESSOS reading a VMEC equilibrium; an ESSOS coil field entering a vmex free-boundary solve goes the other way through from_coils().

  • trace_alphas() — trace fusion-born alpha particles launched from one flux surface of an equilibrium and return the exact loss-fraction diagnostics as an AlphaTracingResult.

The tracer is ESSOS (essos.dynamics.Tracing over essos.fields.Vmec), imported inside the call so vmex imports without ESSOS installed. Only the released ESSOS surface is used — Vmec(wout_filename), Particles, Tracing and its loss_fractions/lost_times/trajectories outputs. Consequences of that restriction:

  • essos.fields.Vmec reads a wout file, so an in-memory equilibrium (WoutData) takes a temporary-wout hop through write_wout(). The file write severs any gradient; the differentiable loss-fraction objective is a separate feature gated on the ESSOS array constructor (uwplasma/ESSOS#61) and is not provided here.

  • The particle energy along each orbit is reconstructed locally from 0.5*m*v_par^2 + mu*|B| (released ESSOS exposes it inconsistently across versions: an eager array in 0.16, a method on the development line).

  • Released ESSOS keeps no solver-failure ledger, so particles_failed counts trajectories with non-finite samples that ESSOS did not attribute to a boundary loss; axis terminations are reported when the installed ESSOS tracks them and are zero otherwise (0.16 has no axis event).

Particles are sampled uniformly on the surface s: theta over [0, 2*pi), phi over one field period [0, 2*pi/nfp), and pitch v_par/v over [-1, 1), from jax.random.PRNGKey(seed). A particle counts as lost when its orbit reaches s >= 0.99 (the ESSOS loss_fraction criterion).

class vmex.core.tracing.AlphaTracingResult(nparticles: int, loss_fraction: float, particles_lost: int, particles_unresolved: int, particles_failed: int, wall_time_s: float, particle_energy: float, total_speed: float, times: ndarray, loss_fractions: ndarray, lost_times: ndarray, trajectories: ndarray, trajectories_xyz: ndarray, energies: ndarray)

Exact alpha-loss diagnostics from one trace_alphas() call.

trajectories holds the guiding-centre coordinates (s, theta, phi, v_par) at each saved time; lost orbits keep the non-finite post-event samples ESSOS writes after the boundary crossing. lost_times is -1 for particles that were never lost.

vmex.core.tracing.essos_vmec_field(source: Any, **kwargs: Any) → Any

Return the essos.fields.Vmec field for an equilibrium or wout file.

source is a path to a wout_*.nc file or an in-memory WoutData. Released ESSOS reads a wout file, so an in-memory equilibrium is written to a temporary wout; ESSOS loads every table eagerly in its constructor, so the file is gone by the time the field is returned. That write severs the gradient — this seam is for diagnostics, not for differentiating through ESSOS.

kwargs reach essos.fields.Vmec unchanged (ntheta, nphi, close and range_torus on the released constructor, which set the resolution of the field.surface ESSOS builds alongside the field).

Released ESSOS reads the stellarator-symmetric wout tables only, so an lasym equilibrium is rejected rather than silently half-transferred.

vmex.core.tracing.trace_alphas(source: Any, *, tmax: float = 0.0003, nparticles: int = 200, s: float = 0.25, seed: int = 42, timestep: float = 5e-07, times_to_trace: int = 200, model: str = 'GuidingCenter') → AlphaTracingResult

Trace fusion alphas from surface s of a wout file or equilibrium.

Parameters:

  • source – Path to a wout_*.nc file, or an in-memory WoutData; handed to ESSOS by essos_vmec_field().

  • tmax – Integration horizon [s], integrator step [s], and number of saved samples (uniform in time, including t = 0).

  • timestep – Integration horizon [s], integrator step [s], and number of saved samples (uniform in time, including t = 0).

  • times_to_trace – Integration horizon [s], integrator step [s], and number of saved samples (uniform in time, including t = 0).

  • nparticles – Ensemble size, launch surface, and sampling seed (see module notes).

  • s – Ensemble size, launch surface, and sampling seed (see module notes).

  • seed – Ensemble size, launch surface, and sampling seed (see module notes).

  • model – ESSOS tracing model ("GuidingCenter" by default).