Metadata-Version: 2.4
Name: ntx
Version: 0.2.2
Summary: JAX-native monoenergetic neoclassical transport solver
Project-URL: Homepage, https://github.com/uwplasma/NTX
Project-URL: Repository, https://github.com/uwplasma/NTX
Project-URL: Documentation, https://ntx.readthedocs.io/en/latest/
Project-URL: Changelog, https://github.com/uwplasma/NTX/blob/main/CHANGELOG.md
Author: UW Plasma
License: MIT
License-File: LICENSE
Keywords: jax,neoclassical,plasma-physics,stellarator,transport
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Physics
Requires-Python: >=3.10
Requires-Dist: interpax
Requires-Dist: jax
Requires-Dist: jaxlib
Requires-Dist: numpy
Requires-Dist: rich
Requires-Dist: scipy
Requires-Dist: tomli; python_version < '3.11'
Requires-Dist: typing-extensions
Provides-Extra: dev
Requires-Dist: build; extra == 'dev'
Requires-Dist: coverage[toml]; extra == 'dev'
Requires-Dist: matplotlib; extra == 'dev'
Requires-Dist: mypy; extra == 'dev'
Requires-Dist: pytest; extra == 'dev'
Requires-Dist: pytest-cov; extra == 'dev'
Requires-Dist: ruff; extra == 'dev'
Requires-Dist: twine; extra == 'dev'
Provides-Extra: docs
Requires-Dist: matplotlib; extra == 'docs'
Requires-Dist: myst-parser; extra == 'docs'
Requires-Dist: sphinx; extra == 'docs'
Requires-Dist: sphinx-rtd-theme; extra == 'docs'
Provides-Extra: io
Requires-Dist: f90nml; extra == 'io'
Requires-Dist: h5py; extra == 'io'
Requires-Dist: netcdf4; extra == 'io'
Description-Content-Type: text/markdown

[![Release](https://img.shields.io/github/v/release/uwplasma/NTX)](https://github.com/uwplasma/NTX/releases)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://github.com/uwplasma/NTX/blob/main/LICENSE)
[![Tests](https://github.com/uwplasma/NTX/actions/workflows/tests.yml/badge.svg)](https://github.com/uwplasma/NTX/actions/workflows/tests.yml)
[![Docs](https://readthedocs.org/projects/ntx/badge/?version=latest)](https://ntx.readthedocs.io/en/latest/)
[![Coverage](https://img.shields.io/badge/coverage-measured%20in%20CI-blue)](https://ntx.readthedocs.io/en/latest/testing.html)
[![Python](https://img.shields.io/badge/python-3.10%20%7C%203.11%20%7C%203.12-blue)](https://pypi.org/project/ntx/)

# 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](https://arxiv.org/abs/2510.27513).

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

```bash
pip install ntx
```

For local development:

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

Optional geometry-coupled examples use upstream JAX geometry tools:

```bash
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:

```bash
ntx examples/example_surface.toml
```

This writes `examples/example_surface.npz`. Plot it with:

```bash
python examples/plot_output_npz.py examples/example_surface.npz
```

Use NTX from Python:

```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:

```python
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 compressed `.npz` outputs.

The input schema is documented in [docs/input-file.md](docs/input-file.md).

## Validation Snapshot

Validation claims are tracked in the maintained
[benchmark matrix](docs/benchmark-matrix.md) and
[physics-gate summary](docs/physics-gates.md). The README keeps only the
highest-signal artifacts:

| Solver validation | Fixed-field current comparison |
| --- | --- |
| ![Monoenergetic validation summary](docs/_static/validation_summary.png) | ![Fixed-field SFINCS, Redl, and NTX + NEOPAX bootstrap-current comparison](docs/_static/bootstrap_current_fixed_field_validation.png) |

| Bootstrap-current workflow | Differentiable geometry/current path |
| --- | --- |
| ![NTX + NEOPAX bootstrap-current profile](docs/_static/bootstrap_current_with_neopax.png) | ![Explicit-relaxed boundary current derivative benchmark](docs/_static/explicit_relaxed_boundary_current_derivative_benchmark.png) |

Current promoted validation includes monoenergetic convergence and identities,
the fixed-field Redl comparison on the precise-QS benchmark family, the
fixed-field `NTX+NEOPAX` total-current stress gate, the integrated W7-X
workflow transfer, and prepared derivative agreement against direct
reverse-mode differentiation. The fixed-field current result uses documented
normalization and moment-closure conventions, not fitted bridge constants.
Species-resolved fixed-field closure parity, broader geometry-family studies,
and large-optimization studies remain tracked as stress diagnostics or planned
research lanes in the docs.

Run the local gate summary with:

```bash
python scripts/check_physics_gates.py
```

## Common Workflows

CLI solves:

```bash
ntx examples/sample_dkes.toml
ntx examples/sample_vmec.toml
```

NEOPAX database and bootstrap-current examples:

```bash
python examples/neopax_with_ntx.py
python examples/bootstrap_current_with_neopax.py
python examples/bootstrap_current_from_vmec_or_boozmn.py
```

Autodiff and optimization examples:

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

Performance examples:

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

Full example coverage is in [docs/examples.md](docs/examples.md).

## Current Open Research Lanes

The major open lanes are:

- full geometry-family reproduction on paper-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,
- restoration of implicit-equilibrium sensitivities only after residual
  contraction and Boozer/NTX transport finite-difference agreement pass,
- 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 is in [docs/research-roadmap.md](docs/research-roadmap.md).

## Documentation

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

## Local Checks

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