Installation

Requirements

• Python 3.10+

• numpy, jax + jaxlib, netCDF4, matplotlib, booz_xform_jax (all installed automatically)

From PyPI

pip install vmex

The plain install includes everything needed for solving, plotting, and the Boozer transform — there are no user-facing extras to remember. Verify with:

vmex --doctor
vmex --test

vmex --doctor diagnoses mixed-Python environments (it prints the active interpreter, pip location, package versions, and JAX backend). If an install misbehaves, first check that pip --version and python -m pip --version point at the same Python.

From conda-forge

conda install --channel conda-forge vmex

or, with Pixi, pixi add vmex. The feedstock may lag PyPI.

From source

git clone https://github.com/uwplasma/vmex
cd vmex
pip install -e .          # editable install, recommended for development

Float64 (required)

VMEC’s numerics require double precision. vmex enables JAX x64 mode itself when you use the CLI or the core solver entry points; if you drive JAX directly in your own scripts, set:

export JAX_ENABLE_X64=1

or jax.config.update("jax_enable_x64", True) before solving.

GPU support

GPU-enabled JAX is intentionally not forced by vmex because the right wheel depends on your platform and CUDA/ROCm version. Install the CPU package first, then install JAX for your accelerator following the official JAX installation matrix, e.g.:

pip install -U "jax[cuda12]"

vmex then picks CPU or GPU per solve using a measured device policy — small decks stay on the CPU, large ones move to the GPU. See GPU guidance for the policy, how to pin a backend with JAX_PLATFORMS, and the persistent compilation cache.

Build the documentation locally

pip install ".[docs]"
python -m sphinx -W -j auto -b html docs docs/_build/html

SPHINX_FAST=1 builds only a minimal landing page for quick CI checks.