ntx 0.2.4

JAX-native monoenergetic neoclassical transport solver

NTX

NTX is a JAX-native monoenergetic neoclassical transport solver for stellarator flux surfaces. It solves the Legendre-space formulation described in Javier Escoto's PhD thesis, Fast monoenergetic neoclassical transport coefficients in stellarators.

Use NTX as:

• a command-line solver for file-backed transport calculations,
• a Python/JAX library for scans, autodiff, uncertainty propagation, and optimization,
• a NEOPAX-compatible monoenergetic database builder for bootstrap-current workflows.

Install

pip install ntx

For local development:

pip install -e ".[dev,docs,io]"

Optional geometry-coupled examples use upstream JAX geometry tools:

pip install git+https://github.com/uwplasma/vmec_jax.git
pip install git+https://github.com/uwplasma/booz_xform_jax.git

Quick Start

Run the smallest bundled case:

ntx examples/example_surface.toml --plot

This writes examples/outputs/example_surface.nc plus a PDF summary panel. Choose the output format by filename:

ntx examples/example_surface.toml --output examples/outputs/example_surface.npz --plot
ntx examples/example_surface.toml --output examples/outputs/example_surface.h5 --plot
python examples/plot_output_file.py examples/outputs/example_surface.nc

Use NTX from Python:

from ntx import GridSpec, MonoenergeticCase, example_surface, solve_monoenergetic

surface = example_surface()
grid = GridSpec(n_theta=17, n_zeta=25, n_xi=32)
case = MonoenergeticCase(nu_hat=1e-3, epsi_hat=0.0)

result = solve_monoenergetic(surface, grid, case)
print(result.D11, result.D31, result.D13, result.D33)

For JAX scans:

import jax.numpy as jnp
from ntx import GridSpec, example_surface, solve_monoenergetic_scan

surface = example_surface()
grid = GridSpec(17, 25, 16)
nu_hat = jnp.logspace(-5, -2, 8)

coefficients = solve_monoenergetic_scan(
    surface,
    grid,
    nu_hat,
    epsi_hat=jnp.zeros_like(nu_hat),
)

Outputs

For each monoenergetic case, NTX computes:

• D11, D31, D13, D33, and D33_spitzer,
• residual and Onsager diagnostics,
• resolved electric-field normalization,
• geometry arrays and run metadata in NetCDF, NPZ, or HDF5 outputs.

The input schema is documented in docs/input-file.md.

Physics In One Paragraph

NTX solves the local monoenergetic drift-kinetic equation on one flux surface, keeping parallel streaming, mirror force, radial-electric-field precession, and Lorentz pitch-angle scattering at fixed speed. The unknown non-adiabatic response is projected onto Legendre polynomials in pitch angle, giving the block-tridiagonal system solved by the code. The two right-hand sides are the radial-transport drive and the parallel-flow/bootstrap-current drive; NTX returns the monoenergetic coefficients consumed by profile and NEOPAX workflows. The full equation, ordering, normalizations, and coefficient definitions are in docs/physics.md.

Validation Snapshot

Validation claims are tracked in the maintained benchmark matrix and physics-gate summary. The README keeps only the highest-signal artifacts:

Solver validation	Fixed-field current comparison

Owned finite-beta bootstrap stress	Differentiable geometry/current path

Owned finite-beta source response	Same-grid profile-current pitch audit

Current promoted validation includes monoenergetic convergence and identities, the fixed-field Redl/SFINCS comparison on the precise-QS benchmark family, the scoped fixed-field NTX+NEOPAX total-current stress gate, the integrated W7-X workflow transfer, and prepared derivative agreement against direct reverse-mode differentiation. Current comparisons use documented normalizations and moment-closure conventions; fitted bridge constants are not used in the runtime.

The owned finite-beta stellarator lane now has same-grid SFINCS-JAX input generation, completed coefficient ladders, Redl and NTX+NEOPAX current audits, true-residual-converged RHSMode=1 SFINCS-JAX profile-current outputs, source-channel decompositions, and radial-interpolation diagnostics. That lane is closed as a reduced-closure stress benchmark: the monoenergetic coefficient differences are below 2.1e-2, while the profile-current pitch-resolution audit accepts the high-Nxi even/odd truncation gap of 1.32e-1 under the 1.5e-1 reduced-closure stress tolerance. The full-collision production branch remains a non-shipping feasibility diagnostic. The corrected Boozer path evaluates B00 on normalized radius and converts dB00/d rho to the physical radial derivative with the VMEC minor radius. The detailed reduced-closure interpretation is in docs/validation.md.

Run the local gate summary with:

python scripts/check_physics_gates.py

Common Workflows

CLI solves:

ntx examples/sample_dkes.toml
ntx examples/sample_vmec.toml

Boozer-file backend validation:

python examples/boozmn_same_coordinate_roundtrip_audit.py

The direct Boozer loader is gated on VMEC half-grid boozmn metadata; the finite-beta finalized-wout transfer audit is documented in docs/geometry.md.

NEOPAX database and bootstrap-current examples:

python examples/neopax_with_ntx.py
python examples/build_neopax_scan_from_ertilde.py --help
python examples/owned_geometry_neopax_dataset.py
python examples/owned_finite_beta_bootstrap_comparison.py
python examples/owned_finite_beta_sfincs_jax_profile_current_audit.py
python examples/owned_finite_beta_source_response_profile_audit.py
python examples/bootstrap_current_with_neopax.py
python examples/bootstrap_current_from_vmec_or_boozmn.py

For large finite-beta NEOPAX database exports on CPU, use --scan-batch-size 32 first; it kept the QI hires 25 x 25 x 60 seven-surface example under 10 minutes and about 1.46 GB peak RSS locally. On GPU, leave full-surface batching enabled unless device memory is the limiter. If a CPU-only laptop still runs too slowly, expose multiple JAX CPU devices before launch and shard the per-surface scan, for example XLA_FLAGS=--xla_force_host_platform_device_count=4 python examples/build_neopax_scan_from_ertilde.py ... --device-backend cpu --parallel-devices 4 --scan-batch-size 32. Run python examples/build_neopax_scan_from_ertilde.py --help and confirm both flags are listed; otherwise the local NTX checkout is stale.

The full finite-beta validation sequence, including same-grid SFINCS-JAX input generation and matched-radius closure audits, is documented in docs/validation.md.

Autodiff and optimization examples:

python examples/derivative_audit.py
python examples/explicit_relaxed_boundary_current_derivative_benchmark.py
python examples/bootstrap_current_optimization.py

Performance examples:

python examples/prepared_geometry_reuse_profile.py --preset smoke
python scripts/benchmark_scaling.py --help

Full example coverage is in docs/examples.md.

Non-Blocking Research Lanes

The current release-facing lanes are closed or explicitly scoped as stress diagnostics. Remaining work is intentionally kept as future research and must not be promoted until its acceptance gates pass:

• full geometry-family reproduction on production-resolution W7-X, QI, QA/QH, and additional stellarator-family inputs,
• reusable hidden-symmetry and omnigenous benchmark families,
• broader geometry-control autodiff with direct AD, prepared adjoints, and finite-difference agreement on reusable geometry families,
• promotion of the implicit-equilibrium derivative path only after residual contraction and Boozer/NTX transport finite-difference agreement pass; the current implicit artifact is a closed non-shipping diagnostic,
• additional dedicated GPU nodes with healthy multi-GPU execution and device-memory timelines,
• broader fixed-field NTX+NEOPAX closure transfer, including species-resolved current decomposition and any future default closure, without regressing the integrated W7-X workflow,
• broader profile, uncertainty, and robust-design studies before promoting stellarator-design claims.

The live roadmap and promotion criteria are in docs/research-roadmap.md and docs/benchmark-matrix.md.

Documentation

• Documentation: ntx.readthedocs.io/en/latest/
• Validation: docs/validation.md
• Performance: docs/performance.md
• GPU notes: docs/gpu.md
• NEOPAX bridge: docs/neopax.md
• Source map: docs/source-map.md

Local Checks

python -m ruff check .
python -m mypy src/ntx
python -m pytest -q
python -m sphinx -b html docs docs/_build/html