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: VMEC input decks 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 and --trace alpha-particle losses.
• 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.11, 3.12 and 3.13 are tested. JAX 0.11 needs Python 3.12 or newer, so a 3.11 environment resolves an older JAX; 3.12 or newer is recommended.

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:

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

Installing into an environment that already holds 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 instead of keeping it, and a nightly CI lane installs vmex[all] over a stale environment to keep it so. DESC (desc-opt, which requires jax<0.10) and other packages that cap JAX below 0.10 need their own environment: vmex[turbulence] and vmex[all] require jax>=0.10.1, so pip either upgrades JAX past their cap (and warns) 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 them.

GPU, conda-forge and source installs

VMEX does not force a GPU build of JAX, because the right wheel depends on the platform and the CUDA or ROCm version. Install VMEX, then JAX for the accelerator following the JAX installation guide (for example pip install -U "jax[cuda13]"), and read the VMEX GPU guide before choosing a device. conda install --channel conda-forge vmex installs the core package; its feedstock may lag PyPI, and the extras above come from pip. For development:

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

The installation guide covers float64, WSL2 and dependency details.

Solve, plot and restart

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.nearby --restart wout_my_case.nc

--scale writes *_scaled at ARIES-CS size (a = 1.7044 m, ⟨B⟩ = 5.8646 T); two factors B R scale by hand. --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, faster than a GPU at this size). 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 against SIMPLE and SIMSOPT. The same 1000 ARIES-CS alphas (positions, pitches, energy) were traced for 10 ms by each code on the same 8 CPU cores. Runtimes exclude compilation and field set-up (benchmark, details). The loss fractions agree within 0.6σ.

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

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.

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.

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.

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.

Results

• 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 (record, 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. In the older ns = 201 table (VMEX 0.3.0) VMEC++ was faster than VMEX on the largest 3-D decks it completed. On two RTX A4000 GPUs no shipped deck or problem size ran faster than on the CPU.
• 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 benchmarks/make_readme_figures.py --only convergence).

What you can build

Every figure below is regenerated by the script named beside it, from decks in this repository.

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.

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

The same scalar-pressure force balance, J x B = grad p on nested flux surfaces, solved by minimizing the MHD energy. What changes is the coordinate system. A tokamak or stellarator has one non-periodic coordinate (the flux label s) and two periodic angles. An open mirror uses (s, theta, xi): two non-periodic coordinates, the flux label s in [0, 1] and the axial coordinate xi in [-1, 1] between two end cuts, and one periodic angle theta. So theta keeps Fourier modes, xi gets clamped cubic B-splines instead of toroidal modes, and the end cuts xi = ±1 become boundary conditions: their nested surfaces are prescribed and flux passes through them, while 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, which is the Clebsch form B = grad Psi x grad alpha with field-line label 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.

Running the examples

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

git clone 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

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.

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% and 2.2% beta, and the field lines of the extended field (coils plus the plasma's virtual-casing field) launched outside it. Closed surfaces continue 3.1 cm past the LCFS at 1.0% beta and 1.5 cm at 2.2% beta, against 2.7 cm in vacuum; beyond them the lines open. 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.

Accuracy and optional polishing

A small VMEC FSQR/FSQZ/FSQL means the discrete solve converged; it does not bound the continuous force error J × B − ∇p. Optional polishing (--polish, or polish=True) re-solves that force on a native quintic-spline representation — constrained Gauss–Newton steps with force-driven knot insertion, then one exact-Hessian Newton step — and certifies it with an independent oracle. To request it from a deck, put this line at the top of the INDATA file (VMEC2000 reads it as a comment; VMEX runs the native polish after the solve):

! VMEX: POLISH_FORCE_BALANCE = .TRUE.

It currently applies to axisymmetric fixed-boundary decks with prescribed pressure and iota (NCURR = 0, GAMMA = 0, LASYM = F); any other deck with the directive stops with an input error. Non-axisymmetric polishing is a research lane.

python examples/force_balance_polishing.py (about a minute with a warm compilation cache) polishes input.shaped_tokamak_pressure to an RMS force of 7e-6 of volavgB²/(μ₀ Aminor_p) with projected stationarity 1e-10. Read back from WOUT files on the same 370-surface mesh, the RMS force falls from 2.0e5 to 42 N m⁻³ over the volume, from 9.8e5 to 190 near the axis (ρ < 0.2) and from 3.2e3 to 23 at the edge; the written file reproduces the native certificate. eps_F is bounded above by 2 by construction and saturates in vacuum; read the dimensional metrics with it.

The same two files through vmex --plot: the summary's ⟨|F|⟩/⟨|∇(B²/2μ₀)|⟩ falls from 3.9e-4 to 7.7e-6. See the polishing reference and the validation record.

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.

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. Multi-device kernel and derivative tests do not yet establish a scalable distributed nonlinear equilibrium solve.

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 contributing, citation, license and the archived plan.