NTX

NTX is a JAX-native monoenergetic neoclassical transport solver based on the Legendre-space drift-kinetic formulation in Javier Escoto’s PhD thesis, arXiv:2510.27513.

NTX provides:

  • built-in analytic sample surfaces

  • DKES-style Boozer inputs

  • VMEC wout inputs through vmec_jax

  • Boozer boozmn inputs through booz_xform_jax

  • direct NEOPAX-style scan and HDF5 mapping helpers

  • CPU and GPU execution through the same JAX solver path

Fastest Start

Install the package:

pip install ntx

Run the bundled analytic sample:

ntx examples/example_surface.toml --plot

That solves one monoenergetic case, prints a Rich summary, and writes a NetCDF payload plus a PDF summary panel.

Inspect that output graphically:

python examples/plot_output_file.py examples/outputs/example_surface.nc

Main User Entry Point

ntx input.toml

What The Code Solves

NTX starts from the local monoenergetic drift-kinetic equation on a single flux surface,

\[\left( v_\parallel \mathbf b\cdot\nabla \mathbf v_E\cdot\nabla \dot{\xi}\partial_\xi - C_L \right) f = s,\]

with fixed speed, Lorentz pitch-angle scattering, and two source terms: the radial-transport drive and the parallel-flow/bootstrap-current drive. Projecting that equation onto Legendre polynomials in pitch angle gives the system actually solved by the code,

\[L_k f^{(k-1)} + D_k f^{(k)} + U_k f^{(k+1)} = s^{(k)}\]

on a single stellarator flux surface, then evaluates the monoenergetic coefficients

\[\hat D_{11},\ \hat D_{31},\ \hat D_{13},\ \hat D_{33},\ \hat D_{33,\mathrm{Sp}}.\]

The CLI is tuned for production solves and output inspection. The imported JAX lane is tuned for scans, autodiff, optimization, NEOPAX coupling, and parallel-throughput workflows.

Documentation Guide

Contents