v0.1.0

vmec-jax v0.1.0

First 0.1 milestone: a research-grade, differentiable JAX implementation of VMEC,
with VMEC2000 scalar parity and an end-to-end differentiable free boundary.

Highlights

• Differentiable equilibria. Fixed-boundary gradients via the implicit-function-theorem
  adjoint (finite-difference validated to ~1e-6 relative); free-boundary gradients via
  virtual casing.
• True single-stage optimization (new). The free-boundary B·n residual is now
  differentiable in the plasma boundary as well as the coils, so a single
  jax.value_and_grad co-optimizes boundary shape and coils simultaneously — the
  gradient threads through the implicit adjoint (boundary) and virtual casing (coils) at
  once. See examples/single_stage_simultaneous_opt.py.
• Multigrid solver with VMEC2000 scalar parity (energies, aspect ratio, beta, iota,
  pressure, boundary shape), checked in-repo via lightweight golden digests.
• Optimization. QA / QH / QP / QI targets; ESS spectral scaling (single-call, no
  max-mode continuation loop); block-tridiagonal implicit Jacobian; perturbation warm-start
  (3.7× fewer solver iterations).
• Physics. Redl bootstrap current with a self-consistency loop; ballooning stability
  (COBRA); omnigenity / QI residual; gyrokinetic turbulence proxies; a built-in Boozer
  transform.
• Performance. CPU and GPU; persistent XLA compilation cache (on by default);
  memory controls (remat / chunking / donation). Warm fixed-boundary solves are sub-0.1 s.
• Engineering. < 10 MB repo, ~95 % test coverage, SOLVAX-backed structured linear
  algebra, sharded CI, VMEC2000 parity without stored large wout files.

Install

pip install vmec-jax          # PyPI (recommended)
conda install -c conda-forge vmec-jax

Notes

• Free-boundary / virtual-casing features need the optional virtual_casing_jax dependency.
• Python 3.10–3.12.