NEO_JAX

NEO_JAX is a standalone neoclassical transport code for stellarators.

It computes effective ripple and related trapped-particle diagnostics from Boozer-coordinate geometry, provides a clean Python API, integrates directly with Geometry Inputs and Pipelines, and uses JAX for just-in-time compilation, vectorized surface evaluation, and automatic differentiation.

NEO_JAX solver stack and data flow

Installation

Quickstart

Python API

Applications

What The Documentation Covers

Physics model

Governing equations for \(\epsilon_{\mathrm{eff}}^{3/2}\), trapped-particle classes, field-line integrals, and reference scalings.

Theory
Numerics and algorithms

Fourier reconstruction, spline interpolation, Newton refinement, RK4 integration, rational-surface handling, and workload controls.

Numerics
Geometry and pipelines

boozmn inputs, in-memory Boozer objects, VMEC→Boozer→NEO pipelines, and the exact data model consumed by the solver.

Geometry Inputs and Pipelines
JAX and differentiation

What stays on device, which paths are JIT-friendly, and how NEO_JAX is used in optimization and design studies.

Differentiability
Inputs, knobs, and outputs

NeoConfig, control-file fields, runtime environment variables, CLI switches, diagnostics, and result containers.

Configuration and Runtime Controls
Validation and testing

Regression structure, CI coverage, performance guardrails, and comparisons against established reference cases where appropriate.

Testing and CI