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-parity numerics, writes standard wout_*.nc output, and — unlike the Fortran original — is differentiable (fixed boundary by implicit differentiation, free boundary through the virtual-casing vacuum field) and runs on CPUs and GPUs.

Why VMEX?

• VMEC2000 parity. The solver ports the VMEC2000 algorithms (steepest-descent moment method, 1D radial preconditioner, Richardson time stepping, spectral condensation, NESTOR vacuum solve) constant-for-constant. Benchmark decks converge in the same iteration counts as VMEC2000 and the wout files agree per-variable (see Performance and validation). An optional 2D block preconditioner cuts iterations 2.5–11x on stiff decks while leaving the default path byte-identical.

• Differentiable. Gradients of fixed-boundary equilibrium properties with respect to boundary shape and profile parameters via implicit differentiation of the converged fixed point (vmex.core.implicit) — no finite differences, no iteration unrolling — validated against central finite differences (see Optimization and differentiability), with an O(1)-memory adjoint. Free-boundary equilibria are differentiable end-to-end through the virtual-casing vacuum field (coil / extcur derivatives), finite-difference-validated (vmex.core.freeboundary_diff). A growing objectives library — quasisymmetry, omnigenity, Redl bootstrap, ballooning stability, gyrokinetic turbulence proxies — plugs straight into a least-squares driver with those exact gradients, reaching precise QA in a single 14.5-minute CPU call (Optimization and differentiability).

• Drop-in workflow. The vmec command reads VMEC2000 input.* namelists and VMEC++-style JSON, prints VMEC2000-format iteration output, and writes wout_*.nc files that load unchanged in simsopt and booz_xform.

• Batteries included. Built-in plotting (vmex --plot), Boozer transform (vmex --booz via booz_xform_jax), spline profiles, multigrid with hot restart, free boundary from mgrid files or directly from coils, and typed zero-crash error handling. The shared linear/adjoint solver layer is factored out into SOLVAX.

Quickstart

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

See Quickstart for a full tour, including the Python API and the Boozer-coordinate workflow.

Benchmark-suite runtimes: vmex (cold and warm) versus VMEC2000 and VMEC++. Warm (compiled-cache) solves are the relevant number for optimization loops; see Performance and validation for the full table.

Documentation

Getting started

• Installation
• Quickstart

Tutorials

• Tutorials

Theory and numerics

• Theory and conventions
• Equations and derivations
• Algorithms
• Confinement physics: quasisymmetry, omnigenity, stability
• Architecture
• Mirror geometry

Reference

• API reference
• CLI reference
• Input file reference
• wout file reference

Performance and validation

• Performance and validation
• Parallelization

Developer guide

• Optimization and differentiability
• Objectives library
• Contributing
• References