VMEX¶
VMEX is a clean-room, JAX-native reimplementation of the VMEC2000 ideal-MHD
equilibrium code for stellarators and tokamaks. It solves fixed- and
free-boundary equilibria with VMEC2000-derived numerics, writes standard
wout_*.nc files that load unchanged in simsopt and booz_xform, and — unlike
the Fortran original — differentiates converged fixed-boundary equilibria via
implicit differentiation. It runs on CPUs and GPUs.
pip install vmex
vmex --test # bundled QH case: solve + wout + plots
vmex input.circular_tokamak # run any VMEC input deck
vmex --plot wout_circular_tokamak.nc
The same solve from Python, with an exact gradient at the end:
import jax
import vmex as vj
from vmex.core import implicit
inp = vj.VmecInput.from_file("input.circular_tokamak")
result = vj.solve_multigrid(inp) # converged equilibrium
p0 = implicit.params_from_input(inp) # differentiable parameters
grad = jax.grad(lambda p: implicit.run(inp, p).wb)(p0)
New here? All of VMEX is the whole mental model on one page;
Installation covers CPU/GPU installs and vmex --doctor.
Learn by doing: your first equilibrium, plots and Boozer coordinates, a first gradient, a first optimization.
Task recipes: run on GPU, restart from a previous run, free boundary, optimization campaigns, troubleshooting.
Every CLI flag, input key, wout variable, objective; the VMEC2000 compatibility contract, the capability contract, and the API.
The theory: the variational problem, spectral representation, preconditioners, NESTOR, adjoint gradients and SOLVAX.
Getting support¶
Bug reports, feature requests, and questions all go to
GitHub issues, which offers a
template for each; include the input file and the output of vmex --doctor.
Diagnose a failing run covers non-convergence, NaNs, and device placement
first. Contributions follow Contributing.
Benchmark-suite runtimes: vmex (cold and warm) versus VMEC2000 and a VMEC++. Warm (compiled-cache) solves are the relevant number for optimization loops; the full generated table is Performance and validation.¶