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 AI coding agents, Claude and Codex, working to a brief by a human owner who makes every physics, licence and publication decision.
Status
| milestone | M4 |
| log entries | 49 |
| tests passing | 539 recorded full-mode results, 539 with gated codes disabled |
| source | 15308 lines + 8083 lines of tests |
| tested code | 8721f24 |
| updated | 2026-10-11 |
The aggregate combines fresh affected-file checks (56 focused checks pass in full and Tier-A-only modesC-157) with retained execution reports for unchanged tests. It is not a fresh full-suite run. Counts and source size are regenerated at the displayed code checkpoint.
Where to start
- Primer: no physics background assumed.
- Action plan: current work, dependencies and completion criteria.
- The log, read from the bottom up.
- Technical internals, the claims register and the decisions, for the details.
How the physics fits together
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 requested absorbed 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 | done 3 Oct 2026 | — |
| M3 | Surrogates | neural ensembles with physics built in, running inside the loop | done 10 Oct 2026 | — |
| M4 | Verification | proved properties, calibrated uncertainty, closed-loop checks per surrogate | in progress | certification review |
| M5 | Uncertainty and validation | uncertainty propagated through the coupled simulation; validation | not started | physicist sign-off |
The log
Show all 49 entries, newest first
What’s next
The action plan gives dependencies and completion criteria.
- M4: verification. Properties and tooling are selected (OQ-22–24), property contracts are written, and the saved transport comparisons pass on ITER and fail on MAST-U (entry 42). Exports pass the approved numeric gate. The owner prioritized MAST-U remediation before further proofs and calibration. Matched flux probes and an exact replay are complete (entry 43); the approved training-only label refresh completed and passes local artifact checks (entry 44). Artifact synchronization and downloaded-checksum verification are complete (entry 45). Retraining preparation passes focused checks and a training-only smoke (entry 46); the concrete compute request was approved and executed once. The candidate completion audit and numerical export checks pass (entry 47); the offline comparison retains mixed regressions. Coupled replay preparation now passes its focused checks and training-only smoke (entry 48); the approved solver replay completes, but reference coverage fails (entry 49). Resolve reference admissibility and review the bounded follow-up next. The approved reference diagnostic is consumed (the diagnostic completes in 517 seconds within its 30-minute bound, with 40 finite converged steps, exact historical temperature/density/poloidal-flux-profile and final-mask equality, and 171 passing saved-output integrity checksC-193); negative raw heat diffusivity explains the rejected face (final face 4 at normalized radius 0.16 has valid inputs but raw electron heat diffusivity -0.02792124888 m²/s; negative heat diffusivity rejects that face at 8 of 40 accepted states; reference rejection is 323/1040 accepted-state face records (31.0577%)C-194). Reference physics disposition remains ahead (OQ-32). K-030 remains open.
- Profile Newton separately (K-022), distinguishing compilation, ensemble evaluation and Jacobian cost. Historical speed-up belongs to v1.
- 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.
The corrected original ITER numerical pilot completes with readable native heating and signed-current profiles (the corrected EC reference completes in 83.16 s with readable native profilesC-156). Native accounting and descriptive saved-reference comparisons are retained; the rendered PostScript page is blank. Production EC coupling contracts, grouped sampling and ITER boundary choices remain open.