booz_xform_jax

Boozer coordinate transforms for VMEC and VMEX equilibria, with differentiable JAX kernels and legacy BOOZ_XFORM input and output files.

Installation

pip install booz_xform_jax

For NVIDIA GPUs, install JAX with CUDA support in the same environment. For development, use pip install -e .[dev] in a clone.

Quick start

Use a STELLOPT control file with booz_xform_jax, xbooz_xform_jax or xbooz_xform:

booz_xform_jax in_booz.mycase F

from booz_xform_jax import Booz_xform

bx = Booz_xform()
bx.read_wout("wout_mycase.nc", flux=True)
bx.register_surfaces([0.25, 0.5, 0.75])
bx.run()
bx.write_boozmn("boozmn_mycase.nc")

Use bx.run_jax() or the functional API for compiled transforms and jax.grad, jax.jvp and jax.vmap. Examples cover geometry, resolution scans and optimization.

Capabilities

Accuracy and speed

Matched CLI transforms agree with STELLOPT BOOZ_XFORM to double precision, including nonzero sine coefficients. Times below include process startup and file I/O: Apple M2 CPU, float64, fastest of three fresh processes, compilation cache disabled.

Fortran is faster on the small tokamak cases; JAX is faster on these stellarator cases. JAX process startup takes 0.71 s; full-process peak memory is 224–514 MiB versus 12–71 MiB for Fortran. Measurements include resolution and memory; the memory plot shows all four cases.

Compiled Python transforms amortize their first-call cost:

These synchronized kernel measurements use JAX 0.9.2 and exclude setup and file I/O. CPU and GPU records include derivative timings and source revisions.

Reproduce the CLI comparison with a STELLOPT xbooz_xform executable on PATH:

python tools/readme_compare.py --repeats 3

Documentation

Quickstart · Theory · Inputs and outputs · Numerics · STELLOPT compatibility · API

Citation and license

Cite the Boozer-coordinate and BOOZ_XFORM literature in the bibliography, together with this repository. Licensed under MIT.