DKX

Differentiable neoclassical transport for stellarators and tokamaks, in JAX.

DKX solves SFINCS v3's linearized drift-kinetic equation on a flux surface for fluxes, flows, bootstrap current, transport matrices and the ambipolar E_r, on CPU or GPU. Every output is differentiable; SFINCS decks run unchanged.

W7-X from its VMEC equilibrium; E_r beside Pablant et al. (2018) (tools/benchmarks/readme_showcase_w7x.py).

Inputs and outputs

A DKX calculation is one case: a versioned document (schema = 1) in TOML, JSON or a Python dict. dkx.Case validates it and hashes it to a case_id. Profiles hold one value per surface.

dkx schema prints every field with its default; dkx validate case.toml checks a file. A run returns a Result, saved as NetCDF, whose arrays carry units in their names: particle_flux_m2_s and heat_flux_W_m2 (surface × species), parallel_current_A_T_m2 (⟨j·B⟩), electric_field_kV_m, the equation residual per surface and, for ambipolar runs, every root (ambipolar_root_kV_m). SFINCS v3 input.namelist decks are the second route, for tangential drifts, Phi1 and transport matrices: dkx input.namelist writes SFINCS's sfincsOutput.h5. Inputs · outputs.

At a glance

Install

pip install dkx                    # CPU
pip install -U "jax[cuda12]"       # add for NVIDIA GPUs
dkx doctor                         # checks the environment

Python ≥ 3.11. From source: pip install -e . in a clone. Optional Boozer transforms: pip install -e ".[booz]".

Quickstart

import dkx

case = dkx.Case.from_mapping({           # analytic tokamak, teaching grid: seconds, not converged
    "schema": 1, "name": "tokamak", "run": {"workflow": "profile", "progress": False},
    "geometry": {"format": "analytic", "file": "tokamak", "surfaces": [0.16, 0.25, 0.36]},
    "species": [{"name": "deuterium", "charge": 1, "mass_amu": 2.014,
                 "density_m3": [8.0e19, 7.0e19, 6.0e19], "temperature_keV": [1.0, 0.8, 0.6]}],
    "physics": {"model": "full_local", "collisions": "pitch_angle_scattering", "magnetic_drifts": "dkes", "phi1": "off"},
    "electric_field": {"mode": "prescribed", "value_kV_m": 0.0},
    "resolution": {"theta": 9, "zeta": 1, "pitch": 8, "speed": 4}, "solver": {"method": "auto", "relative_tolerance": 1e-8},
})
result = dkx.run(case)
print("solver route:", result.metadata["solver_route"])
print("particle flux:", float(result.arrays["particle_flux_m2_s"][1, 0]))

From a file, then the check to run before quoting a number:

dkx run examples/01_tokamak_profile/case.toml --out result.nc
dkx plot result.nc                      # profile panel; E_r roots when the run was ambipolar
dkx converge examples/01_tokamak_profile/case.toml
dkx input.namelist                      # a SFINCS v3 deck, unchanged

dkx converge exits zero only when refining every axis leaves the observables unchanged.

Differentiate

import jax, jax.numpy as jnp, dkx

case = dkx.Case.from_file("examples/05_ambipolar_profile/case.toml")
problem = dkx.prepare_er_scan(case, surface_index=1)          # geometry, grids, collisions once

def bootstrap_current(er_kv_m):
    scan = dkx.batched_er_scan(problem, er_kv_m, differentiable=True, retain_full_state=True)
    return jnp.sum(scan.moments["FSABjHat"])

j, dj_der = jax.jit(jax.value_and_grad(bootstrap_current))(jnp.array([-0.2, 0.0, 0.2]))

Implicit differentiation: one transposed solve on the primal's factors. Compiled, a gradient costs 1.00–1.11× its primal on a 16,230-unknown deck (differentiation).

Every derivative example is checked against central differences.

What you can compute

Monoenergetic coefficients and transport matrices

RHSMode = 3 gives D11*, D31*, D33* (Beidler et al. 2011); RHSMode = 2 the thermal matrix. Boxes: SFINCS v3. Rung 04_monoenergetic_scan.

Collision operators

Full linearized Fokker–Planck conserves momentum; pitch-angle scattering is faster.

Ambipolar E_r

W7-X program 20160309.010: electron roots in the core, ion roots at the edge; dkx roots prints every classified root. Rung 05_ambipolar_profile.

Phi1 and impurities

Phi1 adds the in-surface potential and its quasineutrality equation. Rung 09_phi1_and_impurities.

Stellarator optimization with a kinetic bootstrap current

VMEX's QA bootstrap example plus dkx.bootstrap.KineticBootstrapMismatch (equilibrium ⟨j·B⟩ against DKX's), traced VMEX → booz_xform_jax → DKX (finite differences agree to 5.7e-5–1.8e-3). On a laptop CPU (30 min, 4.6 GB) the objective falls from 1.78 to 0.0061 and the DKX mismatch from 1.2e-3 to 1.2e-4 (QA_optimization_bootstrap_dkx.py). Pitch-angle scattering does not conserve momentum: its current is 1.5–1.6× Redl's.

Proved against analytic limits

Tests: test_physics_limits.py, test_transport_limits.py, test_shaing_callen.py, test_collision_physics_gates.py.

Verified against SFINCS, MONKES and YANCC

On 38 upstream decks at equal discretization DKX matches SFINCS v3 to solver tolerance: median 4e-6, full Fokker–Planck decks 1e-8. The monoenergetic coefficients agree with MONKES and YANCC within 6 %, D33 within 0.1 %, on three configurations (cross-code benchmarks).

Where the reference falls short

On an HSX deck, released SFINCS v3 and DKX differ by 12–19 % in bootstrap current. SFINCS drops matrix entries below 1e-12, which removes the ion–electron collision coupling of a cold-ion, hot-electron plasma. With that cutoff set to zero, the two codes agree to 7e-11 (SFINCS benchmarks; fix proposed as landreman/sfincs#27).

The HSX-like gap deck (633,604 unknowns, Nx = 16) is open in both codes on a 36 GiB host: two SFINCS routes run out of memory, its GMRES stagnates at 0.9955, and DKX's Krylov route reaches 2.5e-5. DKX assembles its operator exactly from 4,800 products and solves the 66,004-unknown reduction to 1.3e-14.

Where DKX converges, and why

DKX picks its route from the operator's structure (solver routes):

Direct routes are exact. The Krylov restart grows to 1,000 within a memory budget (Nx = 16: 357 iterations against 2,788 at restart 200). Every route reports its true residual.

One factorization, many solves

Pass SolveResult.factors back with factors=; transpose=True solves the adjoint. Three right-hand sides drop from three factorizations to none, 0.96 s to 0.28 s; the sparse adjoint costs 0.15 of a primal (factor reuse).

Performance and memory

HSX_PASCollisions_DKESTrajectories, 744,610 unknowns, against PETSc 3.23 / MUMPS 5.8.2 SFINCS v3.

Cold and warm, M3 Max: 1.72 s and 0.12 s at 40,584 unknowns; 23.6 s and 20.0 s at 744,610. This is one measured 744k-unknown HSX PAS case, where the structured route applies.

Across the upstream suite structured direct is faster on 9 of 9 decks, recycled Krylov on 7 of 23; six decks did not complete. Memory is the weak axis: the JAX floor is about 0.5 GB, and DKX is lighter on 3 of the 32 decks it completed (performance).

Keeping only the Legendre blocks the moments need, 2.53M unknowns fit in 2.21 GB on one RTX A4000.

Examples

Rung 06 shows how far these fast grids are from converged. Gallery: examples.

Documentation, development and citation

Getting started · User guide · Tutorials · Physics · Numerics · Benchmarks · Design decisions · API · Contributing

Reusable solver algorithms live in SOLVAX.

@software{dkx,
  author = {Jorge, Rogerio and contributors},
  title  = {DKX: differentiable drift-kinetic neoclassical transport in JAX},
  url    = {https://github.com/uwplasma/DKX},
  year   = {2026}
}

Please also cite SFINCS (Landreman, Smith, Mollén & Helander, Phys. Plasmas 21, 042503, 2014) when you use its decks or model. Metadata: CITATION.cff. LICENSE