Quickstart

Install directly from PyPI:

pip install vmec-jax

Run the bundled CLI test

The fastest first check after a PyPI install is:

vmec --test

This does not require a source checkout. It copies the packaged input.nfp4_QH_warm_start into vmec_jax_test/, runs the fixed-boundary solver with FTOL_ARRAY = 1e-12 for a faster first check, writes wout_nfp4_QH_warm_start.nc, plots the WOUT file into vmec_jax_test/figures/, and prints the equivalent manual commands.

Run the minimal showcase (recommended)

The simplest way to get started is the axisymmetric showcase. It runs a small suite of input decks, writes a wout_*.nc for each, and produces plots. The optional parity summary uses released VMEC2000 reference wout fixtures that are intentionally not tracked in git. Fetch them first when you want CI-style validation rather than just generating fresh outputs from the inputs:

python tools/fetch_assets.py --list
python tools/fetch_assets.py

For a small free-boundary smoke test that does not require the large asset bundle, use the bundled input.cth_like_free_bdy_lasym_small case together with the tracked mgrid_cth_like_lasym_small.nc file in examples/data.

By default the showcase uses a parity-first single-grid run (--single-ns 13) and VMEC2000-style per-iteration screen output (FSQR/FSQZ/FSQL, RAX, DELT, WMHD):

python -m venv .venv
source .venv/bin/activate
python -m pip install -e .
python examples/showcase_axisym_input_to_wout.py --suite

If you want a release-style non-editable install instead:

python -m pip install .

Run the test suite:

pytest -q

Run the full test suite (requires released netCDF assets):

python tools/fetch_assets.py
RUN_FULL=1 pytest -q

CLI (VMEC2000-style executable)

Once installed (or when working from the repo), you can run vmec like the VMEC2000 executable by pointing it to a single input.* file:

vmec examples/data/input.circular_tokamak

Sanity check (verifies the console script is wired to the right interpreter):

vmec --help

The vmec_jax, vmec-jax, and xvmec_jax command names remain aliases for compatibility. If no console command is found or a command raises ModuleNotFoundError, install and run via the module entrypoint:

python -m pip install -e .
python -m vmec_jax examples/data/input.circular_tokamak

This writes wout_circular_tokamak.nc next to the input file and prints the VMEC2000-style per-iteration screen output by default. Use --quiet to silence the iteration table, and --outdir or --output to control where the wout_*.nc file is written. If you only want a short debug run, pass --max-iter and --no-multigrid (single grid).

Boozer-coordinate CLI workflow

The plain vmec-jax install includes booz_xform_jax. Use vmec --booz to run a Boozer transform after a VMEC solve, or directly from an existing wout_*.nc file. The default transform resolution is mbooz = 32, nbooz = 32, with all VMEC surfaces included:

vmec --booz input.nfp4_QH_warm_start
vmec --booz --plot input.nfp4_QH_warm_start
vmec --booz wout_nfp4_QH_warm_start.nc
vmec --plot boozmn_nfp4_QH_warm_start.nc

--booz --plot writes the usual wout_*.nc, runs booz_xform_jax, writes boozmn_*.nc, and then creates:

• mid-radius and LCFS |B| contour-line plots in Boozer coordinates,

• radial Boozer |B| spectra grouped into QA/axisymmetric, QH, mirror, and non-symmetric mode families,

• an LCFS Fourier spectrum for the largest Boozer modes.

Override the transform resolution or selected surfaces from the CLI:

vmec --booz wout_nfp4_QH_warm_start.nc --mbooz 48 --nbooz 48
vmec --booz wout_nfp4_QH_warm_start.nc --booz-surfaces "0.25,0.5,1.0"

Input decks can carry Boozer defaults in a separate namelist. LBOOZ = F is the safe default used by the example inputs; passing --booz overrides it:

&BOOZ_XFORM_JAX
  LBOOZ = F
  MBOOZ = 32
  NBOOZ = 32
  BOOZ_SURFACES = 'all'
/

vmec_jax writes diagnostic wout files from the last available state even when the requested residual tolerance is not reached. These files preserve the computed geometry, profiles, field diagnostics, and residual traces, and mark the solver status with ier_flag plus vmec_jax_converged__logical__ and vmec_jax_status. Treat vmec_jax_status = nonconverged as a diagnostic checkpoint rather than a validated equilibrium.

Load, save, and inspect a wout file

The wout_*.nc file is the standard VMEC output container. Use vmec_jax.wout.read_wout and vmec_jax.wout.write_wout when you want to inspect or round-trip these files directly:

import numpy as np

import vmec_jax as vj
from vmec_jax.wout import read_wout, write_wout

wout = read_wout("wout_nfp4_QH_warm_start.nc")
write_wout("wout_nfp4_QH_warm_start_copy.nc", wout, overwrite=True)

s = np.linspace(0.0, 1.0, int(wout.ns))
print("aspect =", float(wout.aspect))
print("mean edge iota =", float(wout.iotaf[-1]))

for js, (sj, iota_j) in enumerate(zip(s, wout.iotaf)):
    _theta, _zeta, bmag = vj.vmecplot2_bmag_grid(wout, s_index=js)
    print(f"s={sj:.3f}  iota={float(iota_j): .6e}  <|B|>={float(np.mean(bmag)): .6e}")

The complete runnable example is:

python examples/diagnostics/load_save_wout_profiles.py

It creates a wout from examples/data/input.nfp4_QH_warm_start if one does not already exist, saves a round-trip copy, then prints scalar diagnostics, the iota profile, and simple surface-averaged |B| values.

Spline pressure, iota, and current profiles

VMEC input decks can prescribe profiles as polynomial coefficients or as tabulated spline knots. vmec_jax supports the common VMEC forms:

• PMASS_TYPE = "power_series" with AM coefficients, or PMASS_TYPE = "cubic_spline", "akima_spline", or "line_segment" with AM_AUX_S and AM_AUX_F.

• PIOTA_TYPE = "power_series" with AI coefficients, or PIOTA_TYPE = "cubic_spline", "akima_spline", or "line_segment" with AI_AUX_S and AI_AUX_F when NCURR = 0.

• PCURR_TYPE = "power_series" with AC coefficients for \(I'(s)\), PCURR_TYPE = "power_series_i" with AC coefficients for \(I(s)\), or PCURR_TYPE = "cubic_spline_ip" / "cubic_spline_i" (and matching akima_spline or line_segment forms) with AC_AUX_S and AC_AUX_F when NCURR = 1.

The *_AUX_S arrays are knot locations in normalized toroidal flux \(s \in [0,1]\); the matching *_AUX_F arrays are profile values at those knots. A compact pressure/iota spline example is included:

vmec examples/data/input.profile_splines --plot

Generate editable polynomial and spline pressure/current decks side by side with:

python examples/profile_input_examples.py
vmec examples/outputs/profile_inputs/input.profile_polynomial_pressure_current
vmec examples/outputs/profile_inputs/input.profile_spline_pressure_current

For a finite-beta current-spline example, use:

vmec examples/data/input.nfp4_QH_finite_beta

Free-boundary CLI smoke test

For a small bundled free-boundary case, run:

vmec examples/data/input.cth_like_free_bdy_lasym_small

This input references the tracked examples/data/mgrid_cth_like_lasym_small.nc fixture, so it works in a fresh clone without downloading the large asset bundle. The resulting wout_cth_like_free_bdy_lasym_small.nc can be plotted with:

vmec --plot examples/data/wout_cth_like_free_bdy_lasym_small.nc

For ESSOS/direct-coil finite-beta scans and coil-only free-boundary examples, see Free-Boundary Coil Optimization.

If you want to compare the conservative parity track against the optimized fixed-boundary CLI-style controller from Python, run:

python examples/fixed_boundary_driver_tracks.py \
  examples/data/input.circular_tokamak \
  --quiet --json

That example writes two wout files (parity and optimized) unless you pass --no-write-wout, and it prints a short runtime / fsq_total comparison table at the end.

Kernel parity on reference states (solver-free)

To validate intermediate pipeline quantities on reference wout states (no nonlinear solve), run:

python tools/diagnostics/pipeline_parity_summary.py

By default this covers the 4-axisymmetric benchmark suite (circular_tokamak, purely_toroidal_field, shaped_tokamak_pressure, solovev).

Scalar residual parity (fsqr/fsqz/fsql) on reference states

To compare scalar residuals reconstructed from a reference state against wout.fsqr/fsqz/fsql:

python tools/diagnostics/getfsq_parity_cases.py --solve-metric

End-to-end solve snapshot

To run a short fixed-boundary solve and compare a few end-to-end outputs against released references, fetch the optional WOUT fixtures first:

python tools/fetch_assets.py --bundle wout-fixtures

python tools/diagnostics/end_to_end_solve_parity_summary.py --use-input-niter --fast

Drop --fast and increase --max-iter for a full parity snapshot (longer runtime).

External VMEC2000 run (optional)

If you have the VMEC2000 Python extension installed (vmec + mpi4py + netCDF4), you can run VMEC2000 on an input and compare outputs to released references:

python tools/diagnostics/external_vmec_driver_compare.py --case circular_tokamak

A minimal API sketch (recommended)

Most users should start from the small public API in vmec_jax.api:

import vmec_jax.api as vj

run = vj.run_fixed_boundary(
    "examples/data/input.shaped_tokamak_pressure",
    max_iter=10,
    verbose=True,
)
wout_path = "wout_shaped_tokamak_pressure_vmec_jax.nc"
wout = vj.write_wout_from_fixed_boundary_run(
    wout_path,
    run,
    include_fsq=True,
)
boozmn = vj.run_booz_xform(wout_path, mbooz=32, nbooz=32)
vj.plot_boozmn(boozmn, outdir="figures/")

# If you only need an in-memory wout object (no file I/O):
wout_mem = vj.wout_from_fixed_boundary_run(run, include_fsq=True)

wref = vj.read_wout("examples/data/wout_shaped_tokamak_pressure.nc")
print("fsq_total(ref)=", float(wref.fsqr + wref.fsqz + wref.fsql))
print("fsq_total(new)=", float(wout.fsqr + wout.fsqz + wout.fsql))

For free-boundary decks, prefer the explicit entrypoint:

import vmec_jax.api as vj

freeb = vj.run_free_boundary(
    "examples/data/input.cth_like_free_bdy_lasym_small",
    verbose=False,
    use_initial_guess=False,
)
wout_freeb = vj.wout_from_fixed_boundary_run(freeb, include_fsq=True)
print("wb =", float(wout_freeb.wb))
print("wp =", float(wout_freeb.wp))

Use run_fixed_boundary(...) if you deliberately want one driver that accepts either mode. It remains backward compatible and will still dispatch to the free-boundary path when LFREEB = T in the input deck.

Simple optimization example

For a VMEC-JAX-only optimization workflow with explicit SIMSOPT-style objective construction, run:

PYTHONPATH=. JAX_PLATFORMS=cpu python examples/optimization/QH_optimization.py

The script builds a FixedBoundaryVMEC object, constructs objective tuples such as aspect ratio, iota floor, and quasisymmetry residuals, runs least_squares_solve, then shows how to save and plot the resulting equilibrium. The companion examples/optimization/README.md file lists the recommended standalone examples, sweep/rendering tools, and older comparison scripts.