Building the Tokamak Toolkit in public

An integrated tokamak simulation with guarded neural surrogates

Designing, controlling and assessing a fusion reactor means running plasma simulations thousands of times. The physics codes that are accurate enough take hours per run. Neural networks trained on those codes run in milliseconds, but a network returns an answer for any input, including inputs far outside its training data.

The Tokamak Toolkit is an integrated tokamak simulation in which any expensive physics module can be swapped for a neural surrogate, but only behind a guard. At every radius and every time step the surrogate has to show that it is inside the conditions it was trained on, that its ensemble agrees with itself, and that its answer is physically possible. Where it can’t, a trusted physics model takes over, and the handover is recorded. The whole simulation is written in JAX, so it compiles to one program and can be differentiated end to end, and it speaks IMAS, the fusion community’s standard data format, at every boundary.

The code is private for now. What is public is the argument: the method, the checks that define “correct”, every number with where it came from, and the mistakes. How this is built: the code is written by an AI agent, Claude, working to a brief by a human owner who makes every physics, licence and publication decision.

Status

milestone M2
log entries 17
tests passing 215 full suite, 215 with gated codes disabled
source 7531 lines + 3506 lines of tests
last commit e2fb303
updated 2026-10-03

Where to start

How the physics fits together

How the physics codes in the Tokamak Toolkit exchange information Diagram. TORAX, the core transport solver, sits in the centre and evolves electron and ion temperature, electron density and poloidal flux. Above it, the equilibrium: a prescribed CHEASE EQDSK for the ITER hybrid scenario supplies flux-surface geometry every step; FreeGSNKE computes MAST-U-like free-boundary equilibria from coil currents and its two-way coupling to TORAX is planned. On the left, the sources: TORAX sends density, temperature and Z_eff; RF heating (Gaussian deposition), ASCOT5 fusion alphas (with the magnetic field from the EQDSK, updated outside the loop and held between updates) and other TORAX source models return heat, particle and current sources. On the right, the transport coefficients: TORAX sends gradients, q, magnetic shear, temperature ratio and collisionality; QLKNN, a neural network trained offline on QuaLiKiz, and the Bohm/gyro-Bohm fallback both feed the guard, which picks one per face; with the Sauter neoclassical model they return diffusivities, convection, bootstrap current and conductivity. Below, the pedestal and sawtooth models exchange pressure and q for edge values and profile resets, and TORAX writes profiles, sources and fluxes to IMAS IDSs. Equilibrium Heating, current and particle sources Transport coefficients Turbulence, behind the guard Prescribed equilibrium CHEASE EQDSK file ITER hybrid scenario Free-boundary FreeGSNKE MAST-U-like machine Coil currents Core transport TORAX evolves Te, Ti, ne and ψ on ρ implicit solve every time step Fusion alphas ASCOT5 orbit following, AFSI birth source RF heating Gaussian deposition Other sources TORAX models: auxiliary heating, ohmic, radiation, e-i exchange, gas puff, pellets QuaLiKiz quasilinear code QLKNN neural surrogate Bohm/gyro-Bohm physics fallback Guard per face Neoclassical TORAX, Sauter model Pedestal, sawteeth TORAX models IMAS IDSs flux-surface geometry planned coupling: ↓ ψ(R,Z), geometry ↑ p(ψ), j(ψ), Ip B(R,Z) n, T, Zeff heat, particle, current sources gradients, q, ŝ, Ti/Te, collisionality χe, χi, De, Ve, bootstrap current, conductivity training (offline) pressure, q edge values, profile resets profiles, sources, fluxes every time step, inside the compiled loop outside the loop (ASCOT5 updates, offline training) planned, not yet coupled neural network
How the physics codes exchange information during a run. TORAX solves core transport; the other codes supply its geometry, transport coefficients and sources. Hover over a box for what it does.

Core transport: TORAX. The open-source JAX transport code. On a radial grid in ρ it evolves the electron and ion temperatures, the electron density and the poloidal flux ψ (current diffusion), with an implicit solve every time step (Newton-Raphson, with Picard iteration as a fallback). The Tokamak Toolkit calls TORAX’s own single-step function from a fixed-step compiled loop, so a whole run is one program that can be differentiated end to end. Entries 1 to 3.

Prescribed equilibrium. The CHEASE equilibrium of the ITER hybrid scenario, shipped with TORAX as an EQDSK file. TORAX builds its flux-surface geometry (volumes, areas, metric coefficients) from it. ASCOT5 takes its magnetic field from the same file, with the poloidal flux scaled to TORAX’s plasma current. Used for the ITER cases.

Free-boundary equilibrium: FreeGSNKE. Solves the free-boundary Grad-Shafranov equation for the MAST-U-like machine description and writes an IMAS equilibrium IDS. Not yet coupled to transport: that needs a turbulence model valid on a spherical tokamak, which QLKNN is not. Entry 6.

Coil currents. The poloidal-field coil currents, with the coil, passive-structure and wall geometry of the machine description, are the input to the free-boundary solve.

Fusion alphas: ASCOT5. AFSI computes where D-T fusion alphas are born from TORAX’s density and temperature profiles. ASCOT5 follows each alpha along its guiding-centre orbit while it slows down on electrons and ions, or is lost. The power given to electrons and to ions is mapped back to TORAX’s grid and replaces TORAX’s local alpha-heating model. ASCOT5 runs outside the compiled loop; its result is held as a prescribed source until the next update. Entries 10, 12 and 15.

RF heating. A radial Gaussian deposition of RF power to electrons, normalised to the injected power, with an optional co-located driven current. This is the physics model an RF deposition surrogate would fall back to. Entry 5.

Other sources. TORAX’s own source models: auxiliary heating, ohmic heating, radiation, electron-ion heat exchange, gas puff and pellet fuelling. Which ones are active depends on the scenario configuration.

QuaLiKiz. A quasilinear gyrokinetic turbulence code. QLKNN was trained on its output, offline; the Tokamak Toolkit does not run QuaLiKiz itself.

QLKNN. A neural network that returns turbulent heat and particle fluxes from local plasma parameters (normalised gradients, q, magnetic shear, temperature ratio, collisionality) in microseconds. Entry 4.

Bohm/gyro-Bohm. An empirical transport model. It is used at each face where the guard rejects the surrogate (OQ-1).

The guard. At every radial face and every time step, the surrogate’s output is kept only if its inputs are inside the training domain and its outputs are finite and within bounds. Otherwise the fallback value is used and the substitution is recorded. Entry 4.

Neoclassical. TORAX’s Sauter model: the bootstrap current and the plasma conductivity used in current diffusion.

Pedestal and sawteeth. TORAX’s pedestal model sets the temperature and density at the pedestal top, the outer boundary of the core solve. Its sawtooth model flattens the core profiles inside the q = 1 surface when the trigger condition is met.

IMAS IDSs. The Tokamak Toolkit writes its state as IMAS data structures (core_profiles, core_sources, core_transport, equilibrium). Unit and coordinate conversions happen only in these adapters. FreeGSNKE writes its own equilibrium IDS.

Milestones

Each milestone has acceptance criteria fixed in advance; three of them end in a review by a human expert, not by the agent that built the code.

milestone what it delivers state human gate
M0 Scaffold data structures, the guard, the coupling loop, IMAS adapters, tests done 18 Sep 2026 —
M1 Open baseline open-source physics codes wired in; a published benchmark reproduced done 1 Oct 2026 physicist review: passed
M2 Training data sampled datasets from first-principles turbulence codes, with provenance in progress —
M3 Surrogates neural ensembles with physics built in, running inside the loop not started —
M4 Verification proved properties, calibrated uncertainty, closed-loop checks per surrogate not started certification review
M5 Uncertainty and validation uncertainty propagated through the coupled simulation; validation not started physicist sign-off

The log

Show all 17 entries, newest first

What’s next

  • High-flux cases in the TGLF datasets (K-019). Check a sample with the electromagnetic terms off and against CGYRO, then decide with the physicist on a flux range or a loss weighting before training.
  • Sampling, version 2. A shear-dependent αMHD cap (K-017), joint sampling of the geometry from the measured flux surfaces (K-018), and ITER baseline-scenario shapes (K-016). Then active learning near the critical gradients, as the brief sets out.
  • TGLF inside TORAX (K-015): a transport domain that excludes the axis and the pedestal, so TGLF can be the reference model for checks inside the coupled loop at M3 and M4.
  • ASCOT5 follow-ups, before neutral-beam ions: the alphas’ mean birth energy (K-011) and the core heating profile (K-012).
  • Future extension: extending the benchmark to a STEP power-plant flat-top is planned as a future extension. For now the work stays with the ITER benchmark and MAST-U.
Code and first draft by Claude (Anthropic), working to a brief by Vignesh Gopakumar, who reviewed and approved this page. How this is built · Tokamak Toolkit home