DKX

Differentiable neoclassical transport for stellarators and tokamaks, in JAX.

DKX solves the radially local, linearized drift-kinetic equation on a flux surface and returns particle and heat fluxes, parallel flows, bootstrap current, transport matrices and ambipolar radial electric fields. It implements the SFINCS Fortran v3 model, including full Fokker–Planck collisions and Phi1, reads and writes SFINCS decks, runs on CPU or GPU, and every output is differentiable in every input.

Install

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

Run

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 the shell: dkx run case.toml, dkx converge case.toml (check the resolution before trusting a number), dkx wout_w7x.nc, or dkx input.namelist for a SFINCS deck.

Gradients

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

The derivative passes through the linear solve by the implicit function theorem, and every returned state has satisfied the original kinetic equation. Profiles and geometry differentiate the same way: differentiability.

Choose a workflow

Why DKX

Verified

Against SFINCS Fortran v3 on 38 upstream decks with the same discretization, DKX agrees to solver tolerance: median 4e-6, full Fokker–Planck decks to 1e-8. Against the independent codes MONKES and YANCC, the four Beidler-normalized monoenergetic coefficients agree within 6 percent and D33 within 0.1 percent on three configurations. Gradients agree with central finite differences over a step window on every shipped derivative example. Details, tolerances and scope: validation matrix.

Fast

HSX_PASCollisions_DKESTrajectories, RHSMode=1, 744,610 unknowns, one machine, against the PETSc 3.23 / MUMPS 5.8.2 build of SFINCS v3. Warm is the second solve in a process, after XLA has compiled.

That is one measured 744k-unknown HSX PAS case, chosen because DKX has an exact structured solver for it. Across all 38 upstream decks: structured route faster on 9 of 9; Krylov route faster on 7 of 23, six not completed. Every deck, hardware string and method: performance.

Documentation

Tutorial: first W7-X result · How-to: resolution, wout files, SFINCS decks · Reference: schema, CLI, API · Explanation: physics models, solver routes, limitations

Cite

@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.

Examples · Research plan · Contributing · SOLVAX owns the reusable solver algorithms · LICENSE