NEO_JAX

Effective helical ripple and trapped-particle diagnostics from Boozer geometry, with CPU/GPU kernels, automatic differentiation and a legacy NEO command line.

Installation

pip install neo-jax

The VMEX → Boozer → NEO pipeline requires Python 3.11+, VMEX 0.11.6+ and booz_xform_jax 0.4.3+:

pip install "neo-jax[pipeline]"

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

Quick start

Use an existing neo_in.mycase control file and boozmn_mycase.nc:

neo-jax mycase --quiet

xneo, xneo_jax and python -m neo_jax provide the same interface. The CLI guide covers control-file lookup, progress reporting and legacy output files.

from neo_jax import NeoConfig, run_neo

config = NeoConfig(surfaces=[0.25, 0.5, 0.75], theta_n=64, phi_n=64)
results = run_neo("boozmn_mycase.nc", config=config)
print(results.epsilon_effective)  # epsilon_eff^(3/2), the legacy epstot quantity

Capabilities

File and differentiable pipelines retain both cosine and sine geometry. Float64 is enabled by default; low-|iota| surfaces have an explicit work guard, with exact and approximate policies.

Accuracy and speed

Matched STELLOPT calculations agree in effective ripple across ORBITS and QA surfaces. The validation guide also covers NCSX, geometry, parallel current, CPU/GPU parity and legacy diagnostics.

Apple M2 CPU with JAX 0.10.2; RTX A4000 GPU with JAX 0.9.2; float64 and compilation cache disabled. Fortran includes process startup; JAX includes Boozer file reads and compilation, with imports excluded. ORBITS requests convergence history. Fortran is faster on these cases; measurements record controls, ripple profiles and timing scopes.

Callback-free QA profile: eight surfaces, 388 modes, a 32×32 grid and 24 pitch samples.

Cold includes lowering and compilation. Both exclude transforms, imports, file I/O and plotting; precision, CPU and JAX match the CPU measurements above.

Reproduce one case with a STELLOPT executable:

python benchmarks/benchmark_orbits.py --jax --warmup \
  --control tests/fixtures/orbits/neo_in.ORBITS_FAST \
  --boozmn tests/fixtures/orbits/boozmn_ORBITS_FAST.nc \
  --reference-bin xneo --output comparison.json

Differentiable pipeline

Pass a VMEX input to the host pipeline, or a solved vmex.optimize.Equilibrium to the JAX pipeline for derivatives:

from neo_jax import NeoConfig, run_vmec_boozer_neo

results = run_vmec_boozer_neo(
    "input.mycase", booz_kwargs=dict(mboz=8, nboz=8),
    neo_config=NeoConfig(surfaces=[0.25, 0.5, 0.75], theta_n=32, phi_n=32),
)

Examples cover ripple profiles, geometry derivatives and optimization. The large NCSX fixture is fetched on demand with neo_jax.ncsx_boozmn_path(download=True).

Documentation and citation

Quickstart · Theory · Configuration · Numerics · Differentiability · API · References

Cite the NEO literature in the bibliography, together with this repository. Licensed under MIT.