VMEX

VMEX is a JAX reimplementation of VMEC2000, the standard code for three-dimensional ideal-MHD equilibria of stellarators and tokamaks. It reads VMEC input decks, solves fixed- and free-boundary equilibria (NESTOR vacuum field driven by an MGRID table or coils), reproduces VMEC2000's iteration and writes standard wout_*.nc files. Because the solver is written in JAX, VMEX also differentiates the converged equilibrium through the implicit function theorem, so boundary, profile and coil parameters can be optimized with derivatives of the discrete equilibrium equations.

The goal is fast, differentiable and VMEC2000-compatible equilibria for stellarator optimization, including single-stage design, in which the plasma boundary and the coils are optimized together. VMEX also computes Boozer transforms, fields and their derivatives, quasisymmetry, quasi-isodynamicity and stability diagnostics, and has a separate lane for open mirrors.

• Use existing workflows: VMEC2000 &INDATA decks, VMEC++-style JSON inputs and wout_*.nc output, NS_ARRAY multigrid continuation, hot restart from a saved WOUT, and DESC inputs and outputs read without DESC.
• Design with gradients: implicit scalar adjoints and residual Jacobians for SciPy, JAXopt or Optax, with quasisymmetry, quasi-isodynamic, Mercier, ballooning, bootstrap and maximum-J objectives.
• Inspect the physics: Boozer transforms, the magnetic field and its first three spatial derivatives, effective ripple, the --plot diagnostic summary, --scale to reactor size, --trace alpha-particle losses and --turbulence gyrokinetic runs.
• Choose the hardware: CPU or GPU equilibrium solves (optimization gradients default to CPU), reusable compilation and independent-case ensembles.
• Connect coils: ESSOS coil fields, NESTOR free boundary from an MGRID table or coils, and the virtual-casing exterior field of a finite-beta plasma.

The capability reference defines supported models and validation limits. The toroidal model assumes nested flux surfaces; mirror and high-order polishing features have narrower validation scopes.

Installation

pip install vmex
vmex --doctor
vmex --test

pip install vmex is enough to solve, plot, restart, compute Boozer spectra and optimize with SciPy. Its required dependencies include JAX (CPU), SciPy, netCDF4 and h5py, and two uwplasma packages that the solver itself calls: solvax for the linear solvers (the radial preconditioner, polishing least squares, GMRES) and booz_xform_jax for the Boozer transform. vmex --doctor prints the interpreter, package versions and JAX devices; vmex --test solves a bundled QH deck end to end and writes its figures to ./vmex_test/.

Python 3.10, 3.11, 3.12 and 3.13 are tested. JAX 0.11 needs Python 3.12 or newer, JAX 0.7 and later need 3.11, and SciPy 1.16 needs 3.11, so an older interpreter resolves an older JAX and SciPy; 3.12 or newer is recommended. The coils, neoclassical, turbulence and all extras need Python 3.11 or newer.

Optional packages for the full capabilities

Each extra installs one more package from PyPI and turns on the features that need it. Install all of them with

pip install "vmex[all]"

or pick what you need:

Install	Adds	Enables
pip install "vmex[coils]"	essos>=0.22.0	ESSOS coil fields, vmex --coils free boundary, single-stage plasma and coil optimization, field-line and alpha-particle tracing
pip install "vmex[freeb]"	virtual-casing-jax>=0.0.15	the virtual-casing exterior field of the plasma (VmecExtender)
pip install "vmex[neoclassical]"	neo-jax>=1.1.1, dkx>=2.14.0	effective ripple ε_eff from a WOUT or Boozer spectrum (vmex.epsilon_effective_from_wout), the --plot ripple panel and vmex --neoclassical
pip install "vmex[turbulence]"	gkx>=2.6.0 (with jax>=0.10.1)	gyrokinetic turbulence-proxy objectives (vmex.core.turbulence) and vmex --turbulence
pip install "vmex[optimizers]"	jaxopt, optax	the JAXopt and Optax optimization drivers
pip install "vmex[all]"	all of the above	every example and documented workflow

The same packages can be installed by name; the floors are the ones in pyproject.toml:

Package	Minimum	Installed by	Command
solvax	0.36.2	pip install vmex	pip install "solvax>=0.36.2"
booz_xform_jax	0.5.2	pip install vmex	pip install "booz_xform_jax>=0.5.2"
essos	0.22.0	vmex[coils]	pip install "essos>=0.22.0"
virtual-casing-jax	0.0.15	vmex[freeb]	pip install "virtual-casing-jax>=0.0.15"
neo-jax	1.1.1	vmex[neoclassical]	pip install "neo-jax>=1.1.1"
dkx	2.14.0	vmex[neoclassical]	pip install "dkx>=2.14.0"
gkx	2.6.0	vmex[turbulence]	pip install "gkx>=2.6.0"
jaxopt, optax	none	vmex[optimizers]	pip install jaxopt optax

Existing environments

• Installing over an environment with older packages is supported. Every floor above, and the floors those packages declare (such as solvax's equinox>=0.13.3), upgrades an older copy. A nightly CI lane installs vmex[all] over a stale environment to keep this true.
• DESC (desc-opt) requires jax<0.10, so it needs its own environment. The same holds for other packages that cap JAX below 0.10: vmex[turbulence] and vmex[all] require jax>=0.10.1, so pip either upgrades JAX past their cap (with a warning) or cannot resolve.
• NESTOR free boundary from an MGRID table needs no extra.
• A feature whose package is missing raises an ImportError that names the package to install. The core solver never imports the optional packages.

GPU and conda-forge

• VMEX does not force a GPU build of JAX: the right wheel depends on the platform and the CUDA or ROCm version. Install VMEX, then JAX for your accelerator, for example pip install -U "jax[cuda13]". See the JAX installation guide and the VMEX GPU guide.
• conda install --channel conda-forge vmex installs the core package. The feedstock may lag PyPI, and the extras come from pip.

From source

git clone --filter=blob:none https://github.com/uwplasma/vmex
cd vmex
pip install -e ".[all,dev]"

--filter=blob:none fetches old file versions on demand: the clone is about 8.5 MB instead of 45 MB, and the full history stays available.

JAX settings

• import vmex does not change jax.config. vmex.configure() applies VMEX's settings: float64, a persistent compilation cache, synchronous CPU dispatch and quiet JAX logging.
• The vmex command calls it, and so does the first import of vmex.core or vmex.mirror, unless you called it first or set JAX_ENABLE_X64 yourself.
• The solver needs float64 and stops with a message if it is off.
• In a script that also runs other JAX code, call vmex.configure() once at the top.

More on float64, WSL2 and dependencies: installation guide.

Quick start

With your own VMEC input file:

vmex input.my_case --plot
vmex --plot wout_my_case.nc
vmex --booz wout_my_case.nc
vmex --scale wout_my_case.nc
vmex --trace wout_my_case.nc
vmex input.my_case --turbulence
vmex --neoclassical wout_my_case.nc --nc-profiles profiles.json
vmex input.nearby --restart wout_my_case.nc
vmex wout_my_case.nc --to-input    # writes input.my_case

--scale writes *_scaled at ARIES-CS size (a = 1.7044 m, ⟨B⟩ = 5.8646 T); two factors B R scale by hand.

VMEX follows the deck's NS_ARRAY, FTOL_ARRAY and NITER_ARRAY. In Python:

import vmex as vj

inp = vj.VmecInput.from_file("input.my_case")
result = vj.solve_multigrid(inp, verbose=True)
print(result.converged, result.fsqr, result.fsqz, result.fsql)
vj.write_wout("wout_my_case.nc", vj.wout_from_result(inp, result))

Pass initial_state=result.state for a nearby solve, or restart_from= for a saved WOUT. The restart and CLI guides cover resolution changes, devices, profiles and output controls.

To reconstruct a standalone deck in Python, use vj.VmecInput.from_wout("wout_my_case.nc").to_indata("input.my_case"). The WOUT supplies the final boundary, flux and profiles. VMEX infers NCURR from the profile echo; the original multigrid ladder cannot be recovered. For a free-boundary case, the referenced MGRID file is also needed. If input.my_case exists, the CLI writes input.my_case_from_wout instead.

vmex equilibrium.h5 reads DESC text inputs and HDF5/pickle outputs without installing DESC, using the final stage or equilibrium; it writes input.equilibrium and solves it to write wout_equilibrium.nc. --desc-tol 0 retains all boundary modes. The default 1% boundary tolerance does not guarantee magnetic-field accuracy. WOUT iota has the opposite sign to DESC.

Diagnostic plots

--plot writes five PNGs beside the input or in --outdir: the summary below, flux-surface cross-sections, |B| in VMEC angles, Mercier stability and the 3-D LCFS. The summary adds Boozer |B|, a J map, D_R and the DESC-normalized force balance (effective ripple needs vmex[neoclassical]). The QA and QI panels are vmex examples/data/input.nfp2_QA_finite_beta --plot and vmex examples/data/input.nfp4_QI_finite_beta --plot.

Differentiate and optimize

Compute a boundary/profile derivative without differentiating through every forward iteration:

import jax
from vmex.core import implicit

params = implicit.params_from_input(inp)
gradient = jax.grad(lambda p: implicit.run(inp, p).aspect)(params)

For a simple boundary optimization, supply objectives and let SciPy choose the steps. Each tuple is (function, target, cost_weight):

from scipy.optimize import least_squares
from vmex import optimize as opt

problem = opt.VmecProblem.from_tuples(
    inp, [(opt.aspect_ratio, 4.0, 1.0)], max_mode=1, use_ess=True)
fit = least_squares(
    problem.residual, problem.x0, jac=problem.residual_jac,
    x_scale=problem.scales, max_nfev=20)
problem.input_from_x(fit.x).to_indata("input.optimized")
equilibrium = problem.equilibrium_from_x(fit.x)
vj.write_wout("wout_optimized.nc", equilibrium.wout)

problem.value_and_grad supplies scalar objectives to BFGS/L-BFGS-B, problem.jax_value_and_grad to JAX optimizers, and problem.evaluate(x) reports solve effort and derivative status. The gradient tutorial and optimization guide cover convergence checks, constraints, scaling and finite-difference verification.

axis_field(inp) estimates the axis field before solving; compose it with scale_input to set a target field, or use probe=True for two coarse equilibrium solves.

Ballooning stability

The infinite-n ideal ballooning mode (Dewar and Glasser 1983, Gaur et al. 2023) is the eigenproblem along a field line,

$$ \frac{d}{d\eta}\Big(g,\frac{dX}{d\eta}\Big) + c,X = \lambda,f,X,\qquad X(\pm\eta_b)=0, $$

with line bending g, pressure and curvature drive c and inertia f > 0. A line is unstable when its largest eigenvalue λ = (γ a_N/v_A)² is positive. VMEX discretizes it with finite differences, takes the largest eigenvalue of the resulting tridiagonal matrix with a Richardson step in the grid spacing, and differentiates it with Hellmann–Feynman, dλ = vᵀ dA v, so ballooning_residual works as an implicit least-squares target.

Finite-beta QA seed (input.nfp2_QA_finite_beta, ns = 71) (script, about 3 minutes on a laptop):

• COBRAVMEC parity: with COBRA's three-point radial stencil, four lines agree within 1.7% (tests/test_stability.py). The default five-point stencil reads 2–6% higher and is closer to the converged value.
• Field-line grid: the Richardson step from 128 points per poloidal turn is within 0.05% of the converged eigenvalue; the plain three-point result is 1% low.
• Radial grid: the maximum over 16 field lines and 9 ballooning parameters agrees within 5% between ns = 41 and 101, while λ(s) on one line is rough and changes sign twice between s = 0.80 and 0.83.

Details and the optimization recipe: stability theory and objectives.

What you can build

Every figure below has a generator beside it, using repository decks or retained result data.

Quasisymmetric and quasi-isodynamic boundaries

Boundary optimization for quasi-axisymmetry, quasi-helical symmetry and quasi-poloidal symmetry, each at its own field-period count, from examples/optimization/QA_optimization.py and its QH and QP siblings. python examples/plot_optimized_families.py draws the panel.

Quasi-isodynamic designs at four field-period counts: the |B| contours close poloidally as nfp rises. The QI scripts in examples/optimization/ cover vacuum, finite-beta, bootstrap-consistent and maximum-J variants, with scalar-adjoint, SciPy, JAXopt and Optax drivers.

Quasi-isodynamic design with a self-consistent bootstrap current

The bootstrap current of a quasi-isodynamic field depends strongly on the radial electric field, and the Redl formula, which assumes quasi-axisymmetry, does not describe it (13 to 16 MA T/m² here against 0.3 to 1.7 from DKX). QI_optimization_bootstrap_er.py alternates a Picard update of the applied current profile to the DKX <j·B> at each surface's ambipolar ion-root E_r with QI shape stages at that fixed current (max_mode 1 then 3). It starts from the vacuum QI boundary input.QI_nfp2_initial at β = 1.5 % (needs vmex[neoclassical]; see the neoclassical guide):

python examples/optimization/QI_optimization_bootstrap_er.py

	seed (consistent current)	optimized
constructed-QI residual	5.6e-3	2.0e-3
mirror ratio (limit 0.21)	0.240	0.208
aspect ratio (limit 10)	10.0	9.9
min |ι| (floor 0.5)	0.56	0.59
β (target 1.5 %)	1.50 %	1.50 %
rms(<j·B> equilibrium − DKX) / rms(DKX)	1.8 %	1.7 %

The consistency error is taken on four kinetic surfaces, between which the profile is a monotone cubic through the DKX values; the script asserts it stays below 5 %. The final toroidal current is 269 kA. A full run takes 10 to 12 min on 8 CPU cores of a shared workstation, compilation included, and the CI smoke pass about 5 min (figure script).

Piecewise omnigenity

pwO_optimization.py minimizes the mapped pwO residual, the variance of |B| along contours of a squeezed coordinate map (Liu et al.) whose parameters a, b, c and eta_min are optimized together with the boundary, under shared one-sided limits on aspect ratio, rotational transform, mirror ratio, axis excursion, elongation and edge Jacobian.

From a near-circular seed, modes 1–3 and no magnetic well, the NFP 3 vacuum run lowers the pwO residual from 1.0e-2 to 1.3e-4 and the mean NEO effective ripple ε_eff^3/2 by 95% (6.4e-3 to 3.2e-4), at aspect ratio 7.03 and min |ι| = 0.64. The elongation ceiling of 5 and the |ι| ceiling of 0.9 keep the alpha losses (vmex --trace, 1000 alphas born at s = 0.3, 5 ms, ARIES-CS size) at 0.0%, 0.2% and 0.3% for seed amplitudes 0.05, 0.03 and 0.08, against 5.3% with a ceiling of 9 and no |ι| ceiling. The finite-beta example runs the same ladder from the plain seed and ramps beta with magnetic-well and Mercier limits; the single-stage example continues from the vacuum run; the docs hold the equations and the workflow.

The mapped target field carries zero bootstrap current; the current lives in the unmapped B_max band, and BOOTSTRAP_TERM = "band" flattens |B| across that band. The docs give the equations.

Free boundary with coils

NESTOR free boundary driven by ESSOS coils (or an MGRID table), ramping beta and tracking the Shafranov shift: python examples/free_boundary_essos_coils.py. The exterior field is described below.

Single-stage plasma and coil design

Accepted iterates of single_stage_optimization.py (left, fixed boundary in vacuum) and single_stage_free_boundary_optimization_finite_beta.py (right, free boundary at 0.5% beta).

examples/optimization/single_stage_optimization.py adjusts the plasma boundary and the coils against one weighted objective, solving the equilibrium implicitly at every step; single_stage_free_boundary_optimization.py couples them through a true free-boundary solve. Both print final plasma and coil metrics. Each free-boundary trial is Newton-refined onto the root of the coupled plasma-vacuum residual, where its gradient is exact.

Open mirrors and stellarator-mirror hybrids

VMEX solves the same scalar-pressure force balance, J x B = grad p on nested flux surfaces, by minimizing the MHD energy. Only the coordinates change.

	Tokamak or stellarator	Open mirror
Coordinates	s and two periodic angles	(s, theta, xi)
Radial / axial	s in [0, 1]	s in [0, 1] and xi in [-1, 1] between two end cuts
Periodic angle	two	one (theta, Fourier modes)
Axial basis	toroidal Fourier modes	clamped cubic B-splines in xi

• The end cuts xi = ±1 are boundary conditions: their nested surfaces are prescribed and flux passes through them.
• The free boundary couples the side surface s = 1 to an open exterior vacuum problem.
• The field is divergence-free by construction, sqrt(g) B^theta = I'(s) - d(lambda)/d(xi), sqrt(g) B^xi = Psi'(s) + d(lambda)/d(theta), B^s = 0. This is the Clebsch form B = grad Psi x grad alpha with alpha = theta + lambda/Psi' - (I'/Psi') xi (mirror geometry).
• A stellarator-mirror hybrid closes two straight mirror legs with curved stellarator returns and uses a periodic spline along its length.

vmex examples/data/input.mirror_two_coil_free_boundary --plot   # a &MIRROR deck, writes mout_*.nc

from vmex.mirror import MirrorInput, solve_mirror
solution = solve_mirror(MirrorInput.from_file("examples/data/input.mirror_two_coil_free_boundary"))

You give the boundary (or the coils), flux and pressure; solve_mirror builds the spline grid, initial state and exterior grid.

Benchmarked against the independent Pleiades Green-function code on the two-coil mirror from vacuum to 10% beta: the on-axis field agrees to 7.5e-4 or better on the finer VMEX grid, the difference halves under refinement, and what remains is the size of Pleiades' own grid error (docs/_static/figures/sources/make_mirror_pleiades_figure.py). More mirror and hybrid examples, and which lanes are validated: mirror guide.

Fields, coils and free boundary

The equilibrium exposes Cartesian B() and its first three derivatives, with VJPs, anywhere inside the plasma (set_points_xyz, set_points_flux); it reads the current 10 to 70 times more accurately than the WOUT file. Outside, vj.VmecExtender.from_file("wout_my_case.nc", external_field=coils.B) adds the plasma's virtual-casing field to the coils, accurate to about 1e-12 down to 0.01 minor radii. Coil and MGRID fields enter free-boundary solves as MgridField (trilinear or tricubic).

The Landreman-Paul QA held by its coils as a free boundary at 1.0% beta, and the field lines of the extended field (coils plus the plasma's virtual-casing field) launched outside it. Without net current (left) closed surfaces continue 3.1 cm past the LCFS, then the lines open. A 4 kA toroidal current (middle, and unrolled on the right) lifts the edge iota to 0.514, so iota = 1/2 falls just outside the plasma: an island chain about 0.9 cm wide opens from the LCFS, closed surfaces surround it to 2.9 cm, and the lines open beyond. The figure is built by VMEX from examples/data alone (docs/_static/figures/sources/make_extender_islands_figure.py). The vacuum case and the method are in the fields and coils guide.

near_surface="auto" (the default) switches to a target-graded rule within a few node spacings of the boundary, for B and its derivatives. (a) Error against an exact axis-filament field on the QA boundary at 32, 64 and 96 source samples per period (the estimate flags the first); the direct rule alone fails within 10% of a minor radius. (b) The plasma field at phi = 0. (c) B_vjp against central differences of re-solved equilibria. (d) A mirror (examples/mirror/mirror_exterior_field.py). Script: docs/_static/figures/sources/make_exterior_field_figure.py; parameters: reference, tutorial.

Opt-in solver options

Off by default, so VMEC2000 iteration parity is kept:

• Newton anchoring (vmex.solve_free_boundary_newton) puts a converged free-boundary state on the exact root of the coupled plasma-vacuum residual (how-to).
• Position control (LPOSITION_CONTROL = T or position_control=vj.PositionControl(target=...)) is a PID law on the axis radius that sets a uniform vertical field to hold a current-carrying free-boundary plasma in place. Example: python examples/free_boundary_position_control.py (explanation).
• vmex.relax_free_boundary_constraint continues a converged free-boundary state with a weak spectral constraint, removing the bias that makes the LCFS differ from the coil field by 1e-3 to 1e-2 in |B|. It helps on vacuum and well-resolved decks (|B| error 3-10x lower, force error 4-15x lower); check it against the coil field at low resolution, where it can be worse. result = vmex.solve_free_boundary(inp, external_field=field), then relaxed = vmex.relax_free_boundary_constraint(inp, result, external_field=field). Example: python examples/free_boundary_relax_constraint.py (how-to).
• solve_multigrid(inp, lambda_m0_floor=1e-3) adds a radial preconditioner for the m = 0 lambda force of fixed-boundary stellarators converged to a very small ftol; iterations fall by up to 9% (iteration).

Tracing from the axis to the wall

Orbits do not end at the LCFS. With ESSOS (vmex[coils]), guiding centres move in VMEC coordinates inside, cross the LCFS at a root-found time, and continue in the field outside (coils or mgrid plus the plasma's virtual-casing field, tabulated once) until they strike a wall, come back in, or reach the end time. Below, 3 keV protons in the vacuum Landreman-Paul QA equilibrium, with a wall 3 cm outside the LCFS (examples/vmex_essos_tracing_to_wall.py).

setup = vj.essos_tracing_fields(wout, coils, wall=0.03)     # field, exterior_field, wall
trace = essos.dynamics.Tracing(**setup, model="GuidingCenterAdaptative", particles=particles, maxtime=4e-5)
trace.wall_hits, trace.wall_positions, trace.returns

The equations, the exterior-field options, the interpolation error and the outputs are in the wall-tracing guide.

Alpha-particle losses

--trace (needs vmex[coils]) scales the same way in memory and traces 500 fusion alphas for 10 ms in Boozer coordinates (under a minute on 10 CPU cores; a GPU is faster for large ensembles). It writes the loss fraction and a figure set: loss against time, loss maps on the boundary, and pitch and loss-time distributions. Production runs set --trace-particles N and --trace-tmax T; cost grows as N x T. --trace-birth volume samples the D-T birth profile, and --collisional adds slowing down and pitch-angle scattering (guide).

--trace is compared with SIMPLE and SIMSOPT on matched alphas, and its mode cut and timestep are checked on twenty equilibria (guide).

The ESSOS tracer behind --trace (here with ESSOS's fixed-step Dopri8 defaults; --trace controls the step adaptively) against SIMPLE, FIRM3D, SIMSOPT and DESC on 256 matched 3.5 MeV alphas from s = 0.5 over 10 ms (DESC: the first 32), on 8 Apple M4 cores (record, from uwplasma/ESSOS#162):

Neoclassical transport

A stellarator's confinement is set as much by neoclassical transport as by its shape. Collisions and drifts in the 3-D field fix the radial electric field, the bootstrap current that changes the rotational transform, and a 1/ν loss channel that tokamaks do not have. vmex --neoclassical (needs vmex[neoclassical]) runs the drift-kinetic solver DKX on a solved equilibrium and returns all three:

• the monoenergetic coefficients D11*, D31*, D33* against collisionality;
• the ambipolar root E_r(r);
• the particle and heat fluxes of each species at that root;
• the bootstrap <j·B>, set against VMEX's Redl formula on the same profiles and the equilibrium's own current.

The --nc-profiles file holds n_e, T_e and T_i as polynomials in s (guide). On the bundled QA, β = 2.5 % deck:

vmex wout_LandremanPaul2021_QA_beta2p5_bootstrap.nc --neoclassical \
     --nc-profiles examples/data/kinetic_profiles.LandremanPaul2021_QA_beta2p5_bootstrap.json

It writes *_neoclassical.png and *_neoclassical.h5. The figure below is drawn from that .h5 (script). Runtime: 6.5 min on 4 CPU cores, compilation included.

Turbulence

--turbulence (needs vmex[turbulence], GKX >= 2.6.0) samples the flux tube at s = 0.5, α = 0 and runs GKX on it: a linear k_y scan, the eigenfunction at the fastest-growing k_y, and a short nonlinear ITG simulation (a/L_T = 3, a/L_n = 1, adiabatic electrons). The terminal shows each stage with an estimate of the time left. The run writes the saturated heat flux with its standard error, a summary figure, the field-line geometry and the flux traces. --turbulence-s, -alpha, -ky, -grid, -moments, -tmax, -gradients and --turbulence-kinetic-electrons change the run (guide). The defaults are a few-minute survey and are not converged.

How VMEX works

1. Energy principle. VMEX finds a stationary point of W = ∫ (B²/2μ₀ + p/(γ−1)) dV over nested flux surfaces, where J × B = ∇p. As in VMEC, the unknowns are the surface shapes R, Z and the angle function λ: Fourier series in both angles, finite differences in radius.
2. Spectral force iteration. Each iteration evaluates the Fourier-projected forces and takes a damped second-order Richardson step with VMEC2000's time-step control and restarts, in its order. A radial tridiagonal preconditioner per Fourier mode makes the step effective (a 2-D block preconditioner is opt-in), and NS_ARRAY runs a coarse-to-fine radial multigrid ladder. The solve has converged when FSQR, FSQZ and FSQL fall below FTOL.
3. Free boundary. NESTOR solves the exterior Neumann problem for the vacuum potential with Merkel's Green's-function method, driven by an MGRID table or a coil field; pressure balance across the boundary moves the plasma surface.
4. Implicit derivatives. At a root F(x, p) = 0 of the force residual, dx/dp = −(∂F/∂x)⁻¹ ∂F/∂p. A scalar objective needs one adjoint linear solve for all parameters, and least-squares Jacobians factor the block-tridiagonal radial operator once. The forward iterations are not recorded, so memory does not grow with their number. The derivative assumes that the state is a root (on a fixed boundary VMEX first Newton-refines it to refine_tol = 1e-10), that ∂F/∂x is invertible there, and that perturbed solves stay on the same branch. It is the derivative of the discrete equations, not of the continuum problem.
5. Virtual casing. VmecExtender adds the plasma currents' field, a surface integral over the boundary field, to the coil field outside the plasma; its error grows near the surface.

The explanation pages derive each step.

Force-balance polishing

A small VMEC FSQR/FSQZ/FSQL means the discrete solve converged, not that the continuous force J × B − ∇p is small. Polishing (polish=True or --polish) re-solves that force on a quintic spline and, for the axisymmetric case, certifies it with an independent oracle. To request it from a deck, add one comment line at the top of the INDATA file (VMEC2000 ignores it):

! VMEX: POLISH_FORCE_BALANCE = .TRUE.

Tokamak (certified; NCURR = 0, GAMMA = 0, LASYM = F). python examples/force_balance_polishing.py takes about a minute.

• RMS force over the volume falls from 2.0e5 to 42 N m⁻³, near the axis from 9.8e5 to 190, at the edge from 3.2e3 to 23.
• Projected stationarity is 1e-10, and the written WOUT reproduces the native certificate.
• eps_F is bounded above by 2 by construction and saturates in vacuum; read the dimensional metrics with it.

Stellarator (finite-beta QA, prescribed toroidal current, no certificate yet). python examples/prescribed_current_polishing.py takes 46 s on a laptop (9 s for the solve, 16 s for the polish).

• The RMS force relative to the pressure gradient falls 12× near the axis, 9× in the bulk and 20× at the edge on a held-out grid.
• L-BFGS stops at |grad| = 4.6e-4 after 150 iterations: a forward result, not a converged one.

vmex.core.strong_force.traced_lift is the differentiable form of the lift, for force-balance certificates inside gradient-based objectives.

See the polishing reference and the validation record.

Comparison with MRX

python examples/vmex_mrx_comparison.py relaxes the 2.5 % beta QA deck with MRX, which does not assume nested surfaces, starting from the VMEX ns = 65 field.

Install the MRX equilibrium solver from its source repository; the PyPI package with that name is a different project. The example uses this revision:

pip install "git+https://github.com/ToBlick/mrx@dd1942e9efae05c57d62c02dc3207f87f0704f18"

• Traced iota, enclosed current and pressure agree with VMEX.
• Near the axis VMEX iota still moves with ns: 0.188, 0.176, 0.166 at ρ = 0.1 for ns = 65, 129, 257. These three values do not establish its continuum limit. A matched, independently converged DESC reference has not been established for this deck.
• MRX reconstructs pressure from its relaxed field rather than imposing VMEX's pressure profile. This is a consistency comparison from a shared initial field, not a matched-profile validation.

Validation and speed

• VMEC2000 parity. CI solves six decks against stored VMEC2000 output and requires wb within 1e-7 relative, iota and R, Z harmonics on three surfaces within 1e-5, and iteration counts within ±25%. On six decks never run before (ITER, W7-X, HSX, ARIES-CS, ESTELL, Nührenberg–Zille), the worst relative difference from VMEC2000 was 2.5e-10, and the iteration counts were identical on five and one apart on the sixth (performance records, VMEX 0.8.1).

• Speed. On those decks (Apple M4 CPU, one run each) VMEX took 0.50–1.34 times the VMEC2000 wall time with a warm compilation cache and 0.60–3.13 times with it cleared; the difference is mostly XLA compilation, which a repeated solve in one process skips.

  CPU/GPU performance depends on the problem, software versions and cache state. Measure complete solves and derivatives on your hardware using the performance guidance.

• Derivatives. CI compares adjoint gradients with central finite differences: four Solov'ev gradients to 1e-6 relative and a 3-D li383 boundary gradient to 2e-4. One warm value-plus-Jacobian evaluation of a 48-parameter QA problem took 16.6 s on an Apple M4 (record, VMEX 0.7.0).

• Validation record. The validation record lists every gate, tolerance and benchmark record, with the measured accuracy of each lane.

The NFP=4 QH deck at ns = 51 in all three codes, from a trace recorded in July 2026 (python docs/_static/figures/sources/make_readme_figures.py --only convergence).

Against exact equilibria

Resolution studies and code comparisons cannot see an error that every resolution or every VMEC-type code shares. Landreman (2026) found closed-form 3-D equilibria with finite pressure and nested surfaces, and Issan et al. (2026) broke their stellarator symmetry with one parameter tau. vmex-benchmark-analytical scores VMEX, VMEC2000 and DESC against them.

• The exact solutions exposed two VMEX defects that resolution studies missed, both fixed: the field at the axis (#452) and refinement before derivatives (#453).
• At mid-radius the continuous field gives the current 23–450 times more accurately than the WOUT file; B is right to 1e-5.
• Without stellarator symmetry the full basis adds no error: tau = 0 reproduces the symmetric solution to every digit, and mirror pairs agree to 1e-9.
• With tau != 0 the iota = 2 family is degenerate: VMEX converges to a neighbouring equilibrium that differs by 8e-4 at every resolution. See the validation record.
• Sengupta et al. (2026) provide exact nonsymmetric vacuum fields with nested surfaces. From their boundary and flux alone, the bundled checks recover the transform to 5.8e-5 (paraboloidal) and 5.4e-5 (ellipsoidal, without stellarator symmetry) for ns = 65, mpol = 7, ntor = 14 over 0.2 < s < 0.9. The near-axis error remains 3e-4–1e-3.
• The sheared finite-pressure families remain a failed recovery benchmark for VMEX and VMEC2000; the exact-reference identities and the successful integer-family solves do not validate them.

Performance and parallel execution

JAX compiles each solve once per array structure and reuses the executable for matching shapes, so measure first-call, cache-reload and warm costs separately, including refinement and gradients for optimization. CPU and GPU performance depend on resolution and workload; the dated measurements are in the performance reference and the benchmark records.

Optimizations need no worker settings. Finite-difference probes already use every available CPU, and an exact-Jacobian run measures no faster on 36 cores than on 1, so it stays one process. Independent starts on every core are an advanced option.

vj.parallel.solve_ensemble(inputs, workers=None) solves independent cases concurrently and returns results in input order, each identical to solving that input alone; workers=1 is the serial baseline. Set worker and device budgets with the ensemble guide.

Defaults for CPU, GPU and whole nodes, and when to change them: running on CPU, GPU and whole nodes.

Running the examples

The examples live in the repository, not in the wheel. From a clone:

git clone --filter=blob:none https://github.com/uwplasma/vmex
cd vmex
pip install -e ".[all]"
vmex examples/data/input.circular_tokamak --plot
python examples/take_fixed_boundary_gradients.py

Application	Runnable starting point	Needs
Tokamak or stellarator equilibrium	vmex examples/data/input.circular_tokamak --plot; other decks in examples/data	core
QA, QH, QP or QI boundary design	examples/optimization, including bootstrap, ballooning, Mercier and maximum-J variants	core
JAXopt and Optax drivers	QI_optimization_jaxopt.py, QI_optimization_optax.py	vmex[optimizers]
Asymmetric boundary design	stellarator_asymmetry vacuum and finite-beta scripts	core
Single-stage plasma and coils	single_stage_optimization_QA_QH_QI.py (set FAMILY), single_stage_free_boundary_optimization.py	vmex[coils]
Fields and spatial derivatives	python examples/vmex_get_B_gradB.py	core
Exterior field from coils and plasma	python examples/vmex_get_B_outside_plasma.py	vmex[coils,freeb]
ESSOS coils and a free-boundary beta scan	python examples/free_boundary_essos_coils.py	vmex[coils]
Free boundary from an MGRID table	python examples/free_boundary_mgrid.py	core; first python tools/fetch_assets.py --bundle reference-nc
Field lines outside a finite-beta plasma	python examples/vmex_fieldline_tracing_finite_beta.py	vmex[coils,freeb]
Effective ripple	python examples/epsilon_effective.py	vmex[neoclassical]
Independent-case ensembles	python examples/parallel_ensemble_scan.py	core
Open mirrors	mirror/mirror_fixed_boundary_nonaxisymmetric.py, mirror/mirror_free_boundary_beta_scan.py, mirror/mirror_exterior_field.py	core; the beta scan needs vmex[coils]; the exterior field needs a virtual-casing-jax with virtual_casing_jax.mirror
Research force-balance polishing	python examples/force_balance_polishing.py	core

The optimization scripts expose resolutions, objective weights and iteration budgets near the top. Inspect those settings before a research run; advanced coil examples need the optional dependencies and versions in the ESSOS guide.

Documentation, development and citation

Start with your first equilibrium and your first optimization, then the API, VMEC compatibility and troubleshooting pages. For development, pip install -e ".[dev]" and python tools/preflight.py --static. Report issues with the input deck and vmex --doctor output. See citation, license and the archived plan.