vmec-jax 0.0.10

End-to-end differentiable JAX implementation of VMEC2000 for fixed and free-boundary equilibria.

vmec-jax

End-to-end differentiable JAX implementation of VMEC2000 for fixed-boundary and free-boundary ideal-MHD equilibria.

Install

From PyPI:

pip install vmec-jax

QI optimization uses booz_xform_jax for the differentiable Boozer transform:

pip install "vmec-jax[qi]"

From conda-forge:

pixi add vmec-jax
conda install --channel conda-forge vmec-jax

Developer install from source:

git clone https://github.com/uwplasma/vmec_jax
cd vmec_jax
pip install -e ".[qi]"

Quick Start

Run the solver with the VMEC2000-style CLI:

vmec_jax input.nfp4_QH_warm_start

Plot any wout_*.nc file:

vmec_jax --plot wout_nfp4_QH_warm_start.nc
vmec_jax --plot wout_nfp4_QH_warm_start.nc --outdir figures/

Use the Python API:

import vmec_jax as vj

fixed = vj.run_fixed_boundary("input.nfp4_QH_warm_start")
freeb = vj.run_free_boundary("input.cth_like_free_bdy_lasym_small")
vj.plot_wout("wout_nfp4_QH_warm_start.nc", outdir="figures/")

Backend Selection

vmec_jax follows the selected JAX backend. If CPU-only JAX is installed, runs use CPU. If GPU-enabled JAX is installed and selected, runs use the accelerator; vmec_jax does not silently force those runs back to CPU.

python -c "import jax; print(jax.default_backend()); print(jax.devices())"
JAX_PLATFORMS=cpu vmec_jax input.nfp4_QH_warm_start
JAX_PLATFORM_NAME=gpu vmec_jax input.nfp4_QH_warm_start
JAX_PLATFORMS=cuda vmec_jax input.nfp4_QH_warm_start

From Python, leave solver_device unset to inherit JAX's default backend, or pass solver_device="cpu" / solver_device="gpu" explicitly.

Best Optimization Examples

Editable optimization examples live in examples/optimization/. Start with examples/optimization/README.md for workflow anatomy, then use docs/optimization.rst for the full method guide and docs/optimization_sweep_results.rst for generated sweep tables, figures, and reproduction commands.

The panels below show the current stellarator-symmetric examples used for the README: initial LCFS, final LCFS, objective history, and initial/final Boozer |B| contours on the outer surface. Extended policy discussion, LASYM panels, finite-beta examples, QI seed robustness, failure modes, and full CPU/GPU sweep tables live in the docs.

Target	Backend	Policy	max_mode	ESS	Final J	QI legacy	Mirror	Aspect	Iota	Wall time
QA	CPU	continuation	3	yes	2.33e-04			5.000	0.4200	6.3 min
QH	CPU	continuation	3	yes	9.68e-03			4.999	-1.6595	4.0 min
QP	CPU	continuation	3	no	6.76e-02			5.019	-0.6255	3.7 min
QI	CPU	continuation	3	yes	2.17e-03	2.17e-03	0.211	5.001	-0.5494	11.3 min

Recreate the four best-result panels:

PYTHONPATH=. JAX_PLATFORMS=cpu python examples/optimization/generate_qs_ess_sweep.py --backend-label cpu --solver-device cpu --policy continuation --problems qa --modes 3 --ess on --rerun
PYTHONPATH=. JAX_PLATFORMS=cpu python examples/optimization/generate_qs_ess_sweep.py --backend-label cpu --solver-device cpu --policy continuation --problems qh --modes 3 --ess on --rerun
PYTHONPATH=. JAX_PLATFORMS=cpu python examples/optimization/generate_qs_ess_sweep.py --backend-label cpu --solver-device cpu --policy continuation --problems qp --modes 3 --ess off --rerun
PYTHONPATH=. JAX_PLATFORMS=cpu python examples/optimization/generate_qs_ess_sweep.py --backend-label cpu --solver-device cpu --policy continuation --problems qi --modes 3 --ess on --qi-qp-preseed off --rerun
PYTHONPATH=. python examples/optimization/render_readme_best_optimizations.py

Seed-robust QI coverage is tracked separately below. Common minimal-seed stress tests live in the docs and are intentionally failure-revealing; inspect their generated status, success, crashed, and diagnostic columns before using them as promotion evidence.

Recreate the QI coverage panel:

PYTHONPATH=. JAX_PLATFORMS=cpu VMEC_JAX_QI_RUN_CASE=qi_stel_seed_3127 \
  VMEC_JAX_QI_OUTPUT_DIR=results/qi_opt/ess/qi_stel_seed_3127_current_public_final \
  python examples/optimization/QI_optimization.py
PYTHONPATH=. python examples/optimization/render_qi_readme_cases.py

Performance, Validation, Release

• Performance notes and benchmark caveats: docs/performance.rst
• Validation and VMEC2000 parity status: docs/validation.rst
• Testing and coverage strategy: docs/testing_strategy.rst
• Release checklist and CI gates: docs/release_checklist.rst

CLI Reference

vmec_jax input.*           run the equilibrium solver and write wout_*.nc
vmec_jax --plot wout.nc    generate diagnostic plots
vmec_jax --parity input.*  force the conservative VMEC2000-style loop
vmec_jax --help            show the full option list