drbx 2.0.0

JAX-first edge and scrape-off-layer plasma code with differentiable, hardware-agnostic solver infrastructure.

DRBX

DRBX is a JAX-based, end-to-end differentiable drift-reduced Braginskii (DRB) code for edge and scrape-off-layer (SOL) plasma turbulence — on both closed and open field lines, in axisymmetric (tokamak) and non-axisymmetric (stellarator) geometry via the flux-coordinate-independent (FCI) approach.

Because the whole model is written in JAX, every simulation is jit-compiled, runs on CPU or GPU unchanged, and is differentiable: you can take gradients of any output (a saturated fluctuation energy, a transport level) with respect to any input (a density gradient, an adiabaticity, a diffusivity) through the solver. To our knowledge no other published DRB SOL turbulence code is differentiable, and none combines differentiability with FCI stellarator geometry.

Documentation: drbx.readthedocs.io.

Stellarator turbulence in three dimensions

Four-field drift-reduced turbulence on a rotating-ellipse stellarator — a torus whose elliptical cross-section rotates with the toroidal angle. The cutaway shows the density fluctuations on a flux surface and through the interior; every frame is a jit-compiled, differentiable JAX step:

Reproduce with examples/stellarator/stellarator_3d_render.py.

The same geometry supports closed and open field lines: core field lines (blue) stay on flux surfaces, while beyond a toroidal limiter the scrape-off-layer field lines (red) end on the limiter plate, where a Bohm sheath drains the plasma:

Reproduce with examples/stellarator/stellarator_3d_render.py.

Install

pip install drbx          # from PyPI
# or, from source:
git clone https://github.com/uwplasma/drbx && cd drbx && pip install -e .

Runtime dependencies are jax, scipy, matplotlib, netCDF4, rich, pillow, and solvax. Python 3.10-3.12.

Quick start

Run a simulation from a TOML deck, or inspect one without running it:

drbx inspect examples/inputs/restartable_diffusion.toml   # resolve and print the plan
drbx run     examples/inputs/restartable_diffusion.toml   # run and write artifacts

From Python, a differentiable turbulence run is a few lines:

import jax.numpy as jnp
import numpy as np
from drbx.native.hasegawa_wakatani import HasegawaWakataniParameters, hw_grid, hw_run

grid = hw_grid(64, 2 * jnp.pi * 8)
params = HasegawaWakataniParameters(adiabaticity=1.0, gradient=1.0)
rng = np.random.default_rng(0)
zeta0 = jnp.fft.fft2(jnp.asarray(1e-2 * rng.standard_normal((64, 64))))
n0 = jnp.fft.fft2(jnp.asarray(1e-2 * rng.standard_normal((64, 64))))
zeta, n = hw_run(zeta0, n0, grid, params, dt=5e-3, steps=500)  # jit-compiled, differentiable

Every example below is a flat script: parameters at the top, run, plot.

Highlights

Turbulence on closed and open field lines. The same multi-mode seed on the rotating-ellipse stellarator, with all field lines closed (top) and with a limiter opening the outer flux surfaces into a sheath-drained scrape-off layer (bottom). Four toroidal cross-sections; the mode pattern differs plane by plane because the flux surfaces rotate:

Reproduce with examples/stellarator/stellarator_turbulence.py.

Island divertor. A sheared rotational transform with resonant perturbations forms island chains and a stochastic edge. The open scrape-off layer emerges from the field itself: multi-transit field-line tracing marks the finite connection-length region, and the turbulence drains through it:

Reproduce with examples/stellarator/island_divertor.py.

Imported fields: real coils and VMEC equilibria. The same closed/open field-line machinery runs on imported fields: the vacuum Biot-Savart field of the Landreman-Paul quasi-axisymmetric coil set (via ESSOS) shows nested closed surfaces inside a chaotic open edge, and a VMEC equilibrium (via VMEX) provides closed surfaces whose traced rotational transform matches the equilibrium's iotaf profile to ~1e-6 — with the open field lines confined to the thin island/stochastic layer just beyond the last closed vacuum surface, in the island-divertor sense:

Coil field (closed core, open edge)	VMEC equilibrium + coil-field SOL

Reproduce with examples/geometry-3D/essos-field-lines/closed_open_vacuum_poincare.py (left) and examples/geometry-3D/vmex/closed_open_field_lines.py (right).

And the turbulence runs on that same imported field: a four-field interchange simulation on the Landreman-Paul VMEC equilibrium (real metric, |B|, and surface-preserving parallel maps; recovered iota ~ 0.42) with a limiter opening a scrape-off layer outside the last closed surface — the closed core is retained while the open SOL drains through the Bohm sheath:

Reproduce with examples/stellarator/landreman_paul_turbulence.py.

Neutrals and detachment. The open SOL flux tube reaches the two-point Bohm steady state; the hermes-3 neutral model (packaged AMJUEL atomic rates, target recycling) builds the neutral cushion; and with an evolved temperature the SOL detaches — the target cools through 1 eV and the target ion flux rolls over (the SD1D benchmark):

Open SOL flux tube	Recycling SOL with neutrals

Reproduce with examples/sol/open_sol_flux_tube.py (left) and examples/sol/recycling_sol.py (right).

Reproduce with examples/benchmarks/b6_detachment_rollover.py.

Gradient-based optimization through the physics

Because the solver is differentiable end to end, control and design problems become gradient computations. Example: detachment control — find the upstream density that places the divertor target exactly at the 1 eV detachment threshold (Dudson et al., PPCF 61, 065008; Body et al., NME 41, 101819). The sensitivity dTe_target/dn_up is computed by forward-mode autodiff through the entire 20,000-step stiff SOL solve, and a trust-region Newton iteration walks down the detachment cliff to the threshold:

Reproduce with examples/autodiff/detachment_control.py.

A second example: turbulence optimization — find the adiabaticity (the parallel electron conductivity) at which saturated drift-wave transport drops to a quarter of its hydrodynamic-regime level, the classic hydrodynamic-to-adiabatic transition of Hasegawa-Wakatani turbulence (Camargo, Biskamp & Scott, Phys. Plasmas 2, 48 (1995)). A damped Newton iteration on the adiabaticity, with forward-mode gradients through the saturated turbulence, converges in 7 iterations; an independent long run verifies a 3.96x flux reduction against the 4x target. Left column: the initial hydrodynamic state; right column: the optimized adiabatic state:

Reproduce with examples/tokamak/hasegawa_wakatani_optimization.py.

The same machinery recovers a transport-drive parameter by gradient descent through nonlinear drift-wave turbulence (inverse design), and the differentiation-methods example measures which method is cheapest (forward mode for a few parameters, ~2x a forward run; reverse mode for parameter fields; checkpointing when memory-bound).

Performance and parallel execution

Single-CPU turbulence throughput is about 2 million cell-updates per second in float64, and one gradient through a 200-step rollout costs 2-3x a forward run. The full four-field FCI step — four RHS evaluations, each with a GMRES potential inversion — compiles as a single jit program with no host synchronization inside, which roughly halves the stellarator-turbulence step time on one CPU (details):

Turbulence performance	Cost of each differentiation method

Reproduce with examples/benchmarks/performance_benchmark.py (left) and examples/autodiff/differentiation_methods.py (right).

The FCI stack runs across devices with shard_map: the sharded step is bit-exact against single-device execution (checksums identical at every device count), and on a 36-core host with one core per shard a 1.05M-cell step reaches a 4.5x speedup at 16 shards, while one NVIDIA A4000 GPU runs the same step ~21x faster than a single CPU shard. Both ratios are modest precisely because the single-device kernel is now fast — the solvax GMRES potential solve and the sync-free RHS cut the per-step baseline sharply, so there is simply less work left to parallelize (demo):

Reproduce with examples/benchmarks/fci_sharded_strong_scaling.py.

Hasegawa-Wakatani benchmark

The standard two-field drift-wave turbulence benchmark: grown from noise through the linear instability (growth rate verified against the analytic dispersion relation to ~1e-14) into nonlinear E×B transport. Detailed verification figures — dispersion scans, MMS convergence orders, and gradient-vs-finite-difference checks — are in the documentation:

Reproduce with examples/tokamak/drift_wave_turbulence.py.

Reproducing the figures and movies

Every figure and movie above states the script that generates it. Each is a flat pedagogical file in examples/: all parameters at the top with comments, explicit model/geometry/boundary-condition setup through the public API, progress printed while it runs, plot written at the end. The embedded copies live compressed in docs/media/; the scripts regenerate full-quality versions under output/.

What it does

Capability	What ships
Turbulence models	Hasegawa-Wakatani drift-wave (pseudo-spectral, differentiable); FCI 2-field, 4-field interchange (density/vorticity/parallel flows), and electromagnetic drift-reduced stacks with curvature and vorticity/potential closures
Geometry	Rotating-ellipse stellarator (closed core + optional limiter SOL), island-divertor field (emergent stochastic SOL), shifted-torus helical flux tube, open slab SOL, imported ESSOS coil / VMEC / hybrid equilibria — metrics by autodiff of analytic embeddings
Field-line topology	Closed and open field lines; FCI traced field-line maps; multi-transit connection-length tracing; Bohm sheath + target recycling closure on open endpoints
Neutrals (hermes-3 model)	Packaged AMJUEL ionization/recombination + charge-exchange rates (no external database); Galilean-invariant plasma-neutral coupling; recycling SOL and a self-consistent detaching SOL (implicit Spitzer conduction, self-limiting radiation, SD1D rollover)
Linear solver	drbx.linear linearizes any model about an equilibrium; drift-wave, shear-Alfven, and interchange dispersion reproduced to machine precision
Differentiability	jit/grad/vmap through every model — sensitivity, uncertainty propagation, inverse design, detachment control; forward/reverse/checkpointed methods measured and gated to agree
Parallelism	Multi-device shard_map FCI stepping (bit-exact vs single device) with halo exchange; CPU strong scaling demonstrated, GPU-ready
Solvers	Structured solves via solvax (spectral Fourier-Helmholtz elliptic, tridiagonal, Krylov, preconditioners)
Runtime	TOML-deck CLI (drbx inspect / run) and a small Python API; restartable runs; portable JSON/NPZ artifacts

Validation

drbx is validated against a ladder of literature-anchored benchmarks. Each rung has a test (or a documented gate) and an example that regenerates its figure.

Verified today (each with a passing test):

Case	Anchor	What is checked
Method of manufactured solutions	Riva et al., Phys. Plasmas 21, 062301 (2014); Dudson et al. 23, 062303 (2016)	operator / 1D-fluid / FCI convergence order → 2
Resistive drift-wave dispersion	Dudson et al., Comput. Phys. Commun. 180, 1467 (2009)	growth rate and frequency vs analytic dispersion
Shear-Alfvén wave dispersion	Stegmeir et al., Phys. Plasmas 26, 052517 (2019)	phase velocity vs analytic (with electron inertia)
Interchange / Rayleigh-Taylor	curvature-driven flute dispersion	growth rate vs √(gκ)·k_y/k analytic
FCI on non-axisymmetric geometry	Shanahan et al., PPCF 61, 025007 (2019, BSTING)	parallel-operator MMS; differentiable rollout (grad vs FD 6e-11)
Rotating-ellipse (l = 2) FCI	Stegmeir et al., Comput. Phys. Commun. 198, 139 (2016, GRILLIX)	direct & traced-field-line parallel gradient converge at order 2 on a genuinely non-axisymmetric metric; shape-differentiable; a seeded four-field filament generates interchange vorticity on the rotating surfaces
Island-divertor field (B8)	Shanahan et al., J. Plasma Phys. 90 (2024, BSTING); GBS island-divertor studies	sheared-iota island chains + stochastic edge; closed core and finite-connection-length open SOL emerge from multi-transit tracing; turbulence drains through the emergent divertor masks
Open-field-line SOL flux tube	two-point / Bohm-sheath SOL theory (Stangeby, The Plasma Boundary of Magnetic Fusion Devices, 2000)	parallel flow reaches Mach 1 at the targets; target density = half upstream; exact Bohm particle balance and sheath-recycling accounting
Neutrals and recycling (hermes-3 model)	hermes-3: Dudson et al., Comput. Phys. Commun. 296, 108991 (2024); AMJUEL atomic rates	physically-correct ionization/recombination/CX rates; exact plasma↔neutral particle & momentum conservation; neutrals conserve on the 3D closed rotating ellipse and recycle on the open slab
SD1D detachment rollover (B6)	SD1D: Dudson et al., Plasma Phys. Control. Fusion 61, 065008 (2019)	self-consistent SOL (evolved temperature, implicit Spitzer conduction, self-limiting radiation): the target cools through 1 eV into the recombining regime and the target ion flux rolls over as upstream density rises; differentiable
Differentiable inverse design	—	gradient descent through turbulence recovers a drive parameter

Planned rungs (seeded-blob inertial scaling and others) are tracked in the project planning notes; benchmark reports live under docs/ and docs/validation_gallery.md.

Examples

Flagship simulations, by geometry:

	Turbulence flagship	Geometry
Tokamak	drift-wave turbulence (Hasegawa-Wakatani; linear phase B2-verified, differentiable) + inverse design	periodic flux tube
Stellarator	turbulence on closed + open field lines (four-field, limiter SOL, movies) + Landreman-Paul turbulence (four-field on the imported LP VMEC equilibrium, closed core + sheath-drained SOL) + 3D renders (cutaway turbulence movie, field-line topology) + island divertor (B8: Poincare, connection lengths, emergent open SOL) + rotating-ellipse FCI (parallel-operator convergence) + seeded filament + differentiable FCI drift-reduced model	rotating ellipse (closed core + limiter SOL) + shifted-torus helical + imported ESSOS/VMEC
Coils (vacuum)	Landreman-Paul closed + open field lines (ESSOS Biot-Savart, Poincare classification)	imported coil field
VMEC equilibria	closed field lines from a wout file (VMEX import; traced rotational transform matches the equilibrium iotaf profile to ~1e-6) + closed + open field lines (coil field with the VMEC last closed flux surface overlaid)	imported VMEC equilibrium (Landreman-Paul precise QA)
SOL (open)	open SOL flux tube (parallel transport to Bohm-sheath targets; two-point steady state) + recycling SOL (neutrals, ionization/recombination, detachment onset)	open slab flux tube

Open-field-line SOL: open slab flux tube — parallel transport to Bohm-sheath-bounded targets, relaxing to the classic two-point steady state (Mach 1 at the targets, target density half the upstream density), with the FCI sheath/recycling closure on the target plates.

Benchmarks, differentiable, and geometry examples:

• Linear dispersion (B2/B3): examples/benchmarks/linear_dispersion.py reproduces the drift-wave and shear-Alfvén dispersion relations from the linear solver.
• Autodiff: gradient-based detachment control (forward-mode sensitivity through the stiff SOL solve, trust-region Newton onto the 1 eV threshold), inverse design through turbulence (recover a parameter by gradient descent through a nonlinear drift-wave run), choosing the most efficient differentiation method (forward vs reverse vs checkpointed reverse — same gradient, different cost), plus sensitivity, uncertainty, and reduced inverse design.
• Stellarator FCI and imported geometry: examples/geometry-3D/.
• Start with examples/model_selection_guide.py to choose a model family, dimension, and boundary conditions.

The examples are self-contained — no external plasma code is needed to run them. Large figures and movies are hosted in GitHub releases so the checkout stays small.

Geometry and parallelization

The FCI operator and domain-decomposition stack (FciGeometry3D, fci_operators, halo exchange) was contributed by Aiken Xie in PR #3 and is incorporated here. Built on it, the drift-reduced two-field step runs across multiple devices with shard_map: the domain is decomposed into halo-exchanged shards and the sharded RK4 step is bit-exact against the single-device step (checked to ~1e-16 for single-device and forced-four-device runs in tests/test_fci_sharded_2field.py). On a 36-core Linux host with one core bound per shard, a 1.05M-cell step reaches a 4.5x speedup at 16 shards (1.18 s → 0.27 s), and one NVIDIA A4000 GPU runs the same step ~21x faster than a single CPU shard (checksums identical) (strong-scaling script, docs).

Documentation

• Physics and numerics: physics_models.md, equation_to_code_map.md, code_structure.md.
• Performance and differentiability: performance_and_differentiability.md, profiling_runtime.md.
• Validation: validation_gallery.md.
• Testing policy: testing_strategy.md.

Testing

pytest -q -m "not slow"                                   # full fast suite
pytest -q -m "not slow" --cov=drbx --cov-branch        # with coverage

CI runs the full fast suite on Python 3.10–3.12.

Releases

The current development series is described in docs/release_notes_2_0_0.md.

Citing

If you use DRBX in published work, please cite this repository (https://github.com/uwplasma/DRBX).

License

MIT — see LICENSE.