3. Inspect a physics module
Select a map node or use Previous / Next. Equations describe the model; implementation status describes the selected machine.
Tokamak Toolkit
Integrated modelling for MAST-U and ITER
A tokamak confines a hot, electrically charged gas called plasma using magnetic fields. To calculate how that plasma changes, we need models for its magnetic shape, heating, transport of heat and particles, and edge conditions. Each model depends on quantities calculated by the others.
Choose a machine, establish the scenario assumptions, then follow a plasma update. Each module explains its role, inputs, outputs, governing relation and implementation limits. The evidence section separates numerical checks from physical validation.
MAST here refers to the MAST Upgrade (MAST-U) modelling route. The drawings are schematic. The MAST-U cases use simulated MAST-U-like geometry; the ITER cases are modelling benchmarks.
How the models fit together
1. Establish the scenario
Notation, units and model names
These are toolkit interface conventions. Backend adapters must convert native units and coordinates before exchanging quantities.
| Quantity | Meaning and location | Unit |
|---|---|---|
| ρ | Normalised toroidal-flux radius; axis 0, boundary 1. Not a geometric distance. | Dimensionless |
| Te, Ti; ne | Electron/ion temperatures and electron density, on cell centres. | eV; m⁻³ |
| ψ; q; V′ | Poloidal flux, safety factor, and dV/dρ from the geometry interface. | Wb; dimensionless; m³ |
| Qe, Qi; sn | Volumetric heating and particle source, on cell centres. | W m⁻³; m⁻³ s⁻¹ |
| χe, χi, De; ve | Heat/particle diffusivities and particle convection, on cell faces. Negative convection is inward. | m² s⁻¹; m s⁻¹ |
| uⁿ; u⁽ᵏ⁾; Δt | Accepted state at time index n; trial at nonlinear iteration k; proposed time step. | State variables; s for Δt |
Model names: TORAX advances core profiles; CHEASE supplies the ITER equilibrium; FreeGSNKE computes the MAST-U-like equilibrium; IMAS provides the data interface. TGLF supplies turbulence responses used to train a surrogate, EPEDNN supplies a learned pedestal response, and JOREK supports the offline stability branch. EC means electron-cyclotron heating and current drive.
Full notation and technical guide2. Follow a coupled time step
Describe the plasmaShape, temperature, density
→Evaluate the modelsHeating, losses, edge conditions
→Update the profilesUpdated plasma, balance, model limits
The updated plasma becomes the input to the next time step.
Data exchange during a time step
Data in transit
Accepted plasma
Trial profiles
An explanatory sequence, not a numerical run or a literal solver call trace. Modules may be reevaluated during each nonlinear solve. Stability remains a separate offline branch.
Select a module to read about its inputs, outputs and current implementation.
● Implemented / measured◐ Limited / incomplete○ Planned / outside scope
4. Understand surrogate selection
A surrogate is a model trained to approximate a more expensive calculation. Before using its prediction, the guard checks the inputs and outputs against specified limits. Change a check below to see when the fallback is used.
Conceptual hard-switch demonstration only; finite values are assumed. These switches do not run a physics model or reproduce a certified domain. Ensemble agreement is not a calibrated uncertainty interval.
5. Read the evidence and its limits
Implementation status is not a validation verdict. Read each result with its recorded scenario, acceptance criterion and claim status.
What the different checks establish
- Numerical convergence
- Refining the time step tests sensitivity to the numerical protocol. It does not establish agreement with an experiment or remove model bias.
- Reference comparison
- Comparing profiles with a declared reference tests accuracy in that case. A failed or non-evaluable comparison cannot support a validation claim.
- Formal local robustness
- A proof bounds a stated property inside a specified input region. An unresolved proof is not a counterexample; an out-of-domain input is outside that contract.
Research status and build log
ITER convergence result
The current continuous baseline passes the stated time-step refinement criterion. all 33 original cases reach five seconds at all three maximum-step settings: 99 of 99 valid trajectories and 264 of 264 passing endpoint comparisonsC-305. the largest changes are 0.219% for coarse-to-middle and 0.0423% for middle-to-fine, both below the 1% acceptance limitC-306.
Read the full result · What changed in the solver and closure · Reproduce the numerical argument. The owner accepts the frozen closure and grants M5 sign-off for the bounded research deliverable (OQ-52). Full current-baseline uncertainty propagation remains incomplete.
The aim of this work is to use learned physics models within a coupled plasma simulation. A surrogate can reduce the cost of repeated calculations, but its accuracy on a test dataset does not tell us how it will behave when the plasma evolves. The coupled calculation changes the inputs the model receives, and errors can feed back through the other modules.
The Tokamak Toolkit connects selected equilibrium, heating, transport and edge models. Guards check the surrogate inputs and outputs against declared limits and record when a fallback is used. Ensemble routes also check model disagreement. The current ITER baseline uses a continuous transport closure near those limits; the technical notes explain how it differs from the earlier hard switch.
The transport calculation uses JAX, with IMAS adapters for exchanging plasma data. Fixed-step paths support differentiation. The adaptive validation checks individual implicit-root derivatives, but does not support reverse-mode differentiation through retry decisions.
The code is private for now. These pages document the methods, results and unresolved problems, with links to the evidence behind each reported number. Methods and authorship describes the use of Claude and Codex for code and drafts, and the owner’s responsibility for physics, licensing and publication decisions.
Status
| milestone | M5 |
| log entries | 91 |
| tests passing | 1177 recorded full-mode results, 1177 with gated codes disabled |
| source | 25603 lines + 15387 lines of tests |
| baseline commit | 6a13bde; continuation source retained with its evidence |
| updated | 2026-11-04 |
Counts aggregate retained execution reports for the displayed code revision, with affected files rerun after code changes. The full suite has 2 skipped tests; the second suite disables gated codes and has 2 skipped tests. Source size is counted from the current checkout; continuation artifacts bind the executed source hashes.
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 original fixed-step compiled loop supports end-to-end differentiation. The new opt-in adaptive validation uses strict Newton and rejected-step rollback; reverse-mode differentiation through its retry decisions is not supported. Entries 1 to 3 and 97 to 98.
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. JOREK uses the same separatrix boundary for its ITER stability reference cases: the EQDSK profiles in case (a), and profiles derived from a saved TORAX state in case (b).
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.
JOREK stability inputs. Case (a) takes the CHEASE EQDSK pressure and FF′. Case (b) uses pressure and current from the saved TORAX ITER hybrid state, on the same EQDSK separatrix. JOREK solves a new equilibrium: case (b) scales FF′ to match the target plasma current, so its safety-factor profile differs from TORAX’s fixed-geometry state. This is an offline snapshot, not a time-step coupling. Input cases.
JOREK input adapter. Converts SI/eV to JOREK’s normalised units and writes profile files and namelists on the host, outside JAX. Physics choices are explicit inputs. Output is read from native HDF5 restart files, rather than an IMAS IDS. Adapter.
JOREK and run wrapper. An external Tier A MHD reference. The implemented ITER workflow solves fixed-boundary equilibria and scans toroidal mode number for linear pedestal stability. The wrapper snapshots inputs and the executable, launches MPI/OpenMP runs, classifies failures and timeouts, and retains outputs with provenance. It is never called inside the compiled transport loop. Run wrapper.
Stability analysis. Fits growth rates from magnetic and kinetic energies, locates the perturbations, and compares time-step and grid checks. The ITER hybrid pedestal modes have passed the time-step check; grid convergence remains open, so the growth rates are not yet accepted under the full convergence criterion. There is no feedback into TORAX and no trained JOREK surrogate; the later role remains a decision for the owner. Checks and limits.
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 | evidence package and guarded model records | package complete; limits retained | certified within stated limits |
| M5 | Uncertainty and validation | ITER prior propagation, sensitivity and reference-label calibration | bounded deliverable closed; full UQ incomplete | reduced-scope sign-off granted |
The log
Show all 91 entries, newest first
No matching items
What’s next
The action plan gives dependencies and completion criteria.
- M4: verification. ITER version 2 is packaged beside the original model. 12 states proved, three unknown and five outside domain under the actual local-robustness toleranceC-326. The records, guard integration and cheap CI re-proof retain the EC-robustness and MAST-U limitations. Exact proof boxes remain separate from ordinary guard margins (all four verified input boxes match the previous loaders; 48 of 48 retained pedestal states keep identical predictions and decisions; none of the 57 EC states is inside the exact proof boxes, so no EC certificate-based output check applies thereC-284). The revised report-only fallback alert policy is adopted under OQ-50; M4 certification is granted within the stated limits under OQ-52.
- JOREK stability (alongside M4). The time-step checks pass, but the doubled-grid checks remain incomplete and the existing resource projection rule stops further links (all three selected fine-grid checks remain incomplete; the measured-cost projection exceeds the existing compute envelope and further submissions stopC-291). No fine-grid growth rate or stability verdict follows.
- M5: the full ensemble numerical gate passes. all 33 original cases reach five seconds at all three maximum-step settings: 99 of 99 valid trajectories and 264 of 264 passing endpoint comparisonsC-305. The successful opt-in method combines the continuous closure, strict solve and rejected-step rollback (the new research baseline uses a continuous guarded closure, Newton trials accepted only after residual decrease, and adaptive rollback of the complete state and auxiliary dataC-304), with explicit physical residual checks (maximum accepted physical-channel L1 is 1.76 × 10−9 against a 10−8 gate, and native mean L1 is 1.01 × 10−14 against 10−10; no crash occursC-307). The new result supersedes the unresolved numerical status for this baseline. Historical prior results retain their original scope. The complete current-baseline driver and design are implemented. three selected five-second trajectories pass the raw audit; only one of ten finite-difference controls survives the derivative timeoutC-327. the 18,432-trajectory design projects roughly 1,710–1,877 node-hours in the measured worker layout, before derivatives and recoveryC-328. OQ-52 accepts the frozen closure and closes M5 as a bounded research deliverable with physicist sign-off. Full UQ remains incomplete; pedestal/equilibrium discrepancy remains unquantified (K-035).
- 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.