Algorithm

NTX solves the monoenergetic drift-kinetic equation in the Legendre basis used in Javier Escoto’s PhD thesis, arXiv:2510.27513.

Core System

The continuous local model is the monoenergetic drift-kinetic equation with parallel streaming, mirror force, radial-electric-field precession, and Lorentz pitch-angle scattering. NTX projects that equation onto Legendre polynomials in pitch angle. For Legendre mode k, the runtime system is

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

The implementation uses:

  • dense block operators in angular space

  • forward Schur recursion over Legendre mode

  • back-substitution of the low-order modes needed for transport coefficients

  • no explicit matrix inverse

The main implementation files are:

  • geometry and angular fields: src/ntx/geometry.py

  • grids and differentiation matrices: src/ntx/grids.py

  • operator blocks: src/ntx/operators.py

  • solve path and scan path: src/ntx/solver.py

  • transport post-processing: src/ntx/transport.py

Surface Families

NTX currently supports:

  • Boozer harmonic surfaces with scalar psi_p, B_theta, B_zeta

  • VMEC harmonic surfaces with covariant and contravariant field components

Surface loaders live in:

  • src/ntx/io.py

  • src/ntx/vmec.py

  • src/ntx/booz.py

  • src/ntx/vmec_jax_backend.py

  • src/ntx/vmec_jax_vmec.py

Differentiable Lane

The imported solve path is JAX-native and supports:

  • jit

  • vmap

  • gradient propagation through scan inputs

  • in-memory NEOPAX array mapping

The CLI and file-writing path is optimized for usability, while the imported path is the one intended for differentiable workflows.