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 vmex

  • 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 solves the normalized local monoenergetic drift-kinetic equation for the non-adiabatic response on a single flux surface,

\[\left[ \xi \frac{1}{B}\left(B^\theta\partial_\theta+B^\zeta\partial_\zeta\right) + \frac{\hat E_\psi}{\mathcal J\langle B^2\rangle} \left(-B_\zeta\partial_\theta+B_\theta\partial_\zeta\right) - \frac{1-\xi^2}{2B^2} \left(B^\theta\partial_\theta B+B^\zeta\partial_\zeta B\right)\partial_\xi - C_L \right] f = s,\]

where \xi=v_\parallel/v, \mathcal J is the Boozer-coordinate Jacobian, and C_L is the Lorentz pitch-angle-scattering operator. Speed is fixed. The two source systems are the radial-transport drive and the parallel-flow drive used by downstream bootstrap-current closures. Projecting this equation onto Legendre polynomials in pitch angle gives the system 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