NTX

A JAX-native monoenergetic neoclassical transport solver for stellarator flux surfaces — and a differentiable one, so a whole design gradient costs one solve instead of one solve per parameter.

pip install ntx
ntx solve --example --nu-hat 1e-2

Why differentiability changes the cost model

Conventional monoenergetic solvers are forward maps: given a surface, they return D11, D31, D13, D33. Sensitivities then cost one re-solve per parameter, which is why neoclassical optimization is usually done with a handful of shape coefficients rather than a realistic boundary.

NTX carries an adjoint through the block-tridiagonal solve, so the gradient with respect to every parameter costs one extra solve — and is exact rather than step-size limited.

Reproduce: python benchmarks/bench_design_derivatives.py --params 1,2,4,8,16,32. The crossover sits near two parameters; below it, finite differences are simpler and perfectly fine.

What else the solver structure buys

import ntx

prepared = ntx.prepare_monoenergetic_system(ntx.example_surface(), ntx.GridSpec(9, 9, 48))
case = ntx.MonoenergeticCase(nu_hat=1e-1, epsi_hat=0.0)

window = ntx.certify_adjoint_window(prepared, case, rtol=1e-6)
print(int(window), window.certified_relative_error)     # 25, 3.2e-07

result = ntx.solve_prepared(prepared, case, adjoint_window=window)

The certificate is honest about its limits. It is a worst-case bound, so it returns a wider window than an oracle would; and on a weakly collisional chain that does not localize it returns the exact window — correct, and no saving. advise_adjoint_window remains the cheap, uncertified estimate.

Physics and scope

NTX solves the local monoenergetic drift-kinetic equation on one flux surface at fixed speed: parallel streaming, mirror force, radial-electric-field precession, and Lorentz pitch-angle scattering. A finite Legendre expansion in pitch angle produces the block-tridiagonal system.

Physics and normalizations · convergence and residual semantics

Validation

Every promoted claim maps to a script, test, committed artifact, acceptance threshold, and documentation entry in the benchmark matrix. Runtime code does not use fitted bridge constants to force agreement with a benchmark.

Monoenergetic convergence and identities	Fixed-field current comparison

The fixed-field result is a scoped reduced-closure comparison, not species-resolved or full-collision parity; assumptions and provenance are in validation. Run the gates with:

python scripts/check_physics_gates.py

Choose A Workflow

Full runnable catalog: docs/examples.md.

Outputs

NetCDF, NPZ, and HDF5, selected by filename suffix, carrying transport coefficients, diagnostics, resolved electric-field normalization, geometry arrays, and run metadata. Schema in docs/input-file.md.

From a source checkout, the bundled TOML writes NetCDF and PDF; choose the format by suffix:

ntx examples/example_surface.toml --plot
ntx examples/example_surface.toml --output result.h5 --plot

Documentation

Getting started · physics · numerics · autodiff · API · examples · validation · glossary · source map

Development

pip install -e ".[dev,docs,io]"
python -m ruff check .
python -m mypy src/ntx
python scripts/test_lane_manifest.py --check
python -m sphinx -W -b html docs docs/_build/html

Optional geometry backends:

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