DKX

Differentiable neoclassical transport for stellarators and tokamaks, in JAX.

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

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

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.

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]))

The same case from a file, and the checks 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 refines every phase-space axis and exits zero only when the observables stopped moving and every returned state satisfies the kinetic equation.

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 through the solve: one transposed solve on the primal's factors. Compiled, a gradient costs 1.00–1.11× its primal on a 16,230-unknown deck (record, #279).

Finite differences cost two solves per parameter; the adjoint costs one.

Every derivative example is checked against central differences (differentiability).

What you can compute

Monoenergetic coefficients and transport matrices

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

Collision operators

Pitch-angle scattering is fast; full linearized Fokker–Planck conserves momentum, which particle-flux coefficients need at high collisionality.

Ambipolar E_r

W7-X program 20160309.010: electron roots in the core, ion roots at the edge. Every root is classified as ion, electron or unstable and kept with its bracketing evaluations; dkx roots prints them. Rung 05_ambipolar_profile.

Phi1 and impurities

Multispecies runs carry 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 self-consistent QA bootstrap example with one added row, dkx.bootstrap.KineticBootstrapMismatch: the mismatch between the equilibrium's ⟨j·B⟩ and the DKX kinetic one, traced VMEX → booz_xform_jax → DKX so VMEX's implicit Jacobian carries it (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). The default pitch-angle-scattering operator 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 with the same discretization DKX agrees with SFINCS v3 to solver tolerance: median 4e-6, full Fokker–Planck decks to 1e-8. The monoenergetic coefficients agree with MONKES and YANCC within 6 %, D33 within 0.1 %, on three configurations (validation matrix).

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 (record; 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 (record). A deck where SFINCS fails and DKX converges is the open target of the solver program.

Where DKX converges, and why

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

Direct routes converge because they are exact. The Krylov route's high-Nx stall was restart stagnation; the default restart grows to 1,000 within a memory budget (Nx = 16: 357 iterations against 2,788 at restart 200). Every route reports the original-equation residual of what it returns.

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 (record).

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. That is one measured 744k-unknown HSX PAS case, chosen because the structured route applies.

Across the upstream suite the route decides the outcome: structured direct is faster on 9 of 9 decks, recycled Krylov on 7 of 23, and six decks did not complete. Memory is the weak axis: the JAX runtime 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, a 2.53M-unknown solve peaks at 2.21 GB on one RTX A4000. On CPU the warm solve is fastest at eight pinned cores.

Examples

Rung 06 shows how far these fast grids are from converged. More: examples/README.md, examples/optimization.

Documentation, development and citation

Tutorial · How-to: case files · CLI · API · Physics and limitations · Solver routes · Validation · Testing · Contributing · Research plan

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