The TGLF surrogate inside the TORAX loop
Two of the M3 acceptance criteria are about the coupled simulation: the surrogate must be differentiable and run inside the M1 loop, and the end-to-end speed-up must be measured. This entry puts the TGLF surrogate of entry 18 into TORAX as a transport model, behind the guard (see the Primer).
How it is wired
At every face and every solver iteration, the TORAX state is turned into TGLF inputs by TORAX’s own TGLF input builder, so the conventions are those of the datasets. The ensemble returns the members’ mean and spread in gyro-Bohm units, and TORAX’s own conversion turns the fluxes into diffusivities and a particle pinch. Where a face’s inputs leave the box of the training inputs, the surrogate returns a non-finite value with finite derivatives; the guard rejects that face and serves the Bohm/gyro-Bohm fallback (OQ-1). The domain is checked in the surrogate’s own input coordinates, as OQ-10 recommended. The whole chain is a JAX function, so TORAX’s Newton solver can differentiate it.
What checks the derivatives: tests compare reverse-mode gradients with central finite differences and with forward mode, for the guarded coefficients and for the core electron temperature after three Newton steps: reverse mode within 10−6 (coefficients) and 10−7 (after 3 Newton steps) of central differences, and within 10−9 and 10−12 of forward modeC-116.
Cost and agreement with TGLF
| one call, 51 faces | RAPTOR L-mode, Newton, 200 steps | |
|---|---|---|
| TGLF | about 103 s (2.0 s per face)C-118 | not run with Newton: TGLF has no derivatives |
| TGLF surrogate | 3.1 msC-117 | 374 s and 405 s (two runs); no stalled step, 57 steps accepted at the coarse toleranceC-120 |
| QLKNN | — | 29 s and 30 s (two runs); 179 of 200 steps stalled, 6 at the coarse toleranceC-121 |
One surrogate call costs a small fraction of a TGLF call on the same faces. End to end the picture is reversed with respect to QLKNN, the neural network TORAX ships: on the M1 benchmark case under Newton the surrogate run takes far longer than the QLKNN run. Every one of the 20 members is evaluated on every face, including the faces the guard hands to the fallback, and Newton differentiates through all of them. What dominates a step has not been profiled; this is ceiling K-022. With the linear solver, and on TORAX’s ITER hybrid case, the run times are much shorter (RAPTOR L-mode, linear solver: surrogate 5.2 s, QLKNN 0.57 s; Newton with the published constant transport 0.35 s, with Bohm/gyro-Bohm 1.6 s. ITER hybrid: surrogate 0.56 s, QLKNN 0.066 s, Bohm/gyro-Bohm 0.030 sC-122).
On the final state of the benchmark run, TGLF and the surrogate were evaluated on the same inputs: Q_e 0.28 (11 faces) and Q_i 0.44 (12 faces); the surrogate is below TGLF on 10 of the 11 and on all 12C-119. This is one state, not a held-out statistic.
The guard served the fallback on a large share of points (0.51 on the RAPTOR L-mode case and 0.27 on the ITER hybrid caseC-123): the training box of the v1 dataset does not cover the axis, the edge geometry or the collisionality of the outer radii on these ITER cases (K-023). Dataset v3 is sampled along TORAX trajectories to close that gap (entry 22), and OQ-20 widened the ITER collisionality range for it. The surrogate’s ion heat flux is the main ion’s, while TORAX’s TGLF model passes the sum over ions (K-024).
The completed comparison uses the linear solver for both models, the same case, fixed steps, post-processing and fallback. Compilation is excluded: TGLF 540.38 s, v1 surrogate 0.780 s, QLKNN 0.126 s; v1 is about 693 times faster than TGLFC-132. Final profile differences from TGLF: T_e 18.52%, T_i 22.09%, n_e 10.25%, psi 0.44%C-133. The runs finish without solver stalls, but this single case does not establish a closed-loop error bound. These timings belong to the earlier v1 model; the final v3 model needs its own matched comparison.
M3: the v1 surrogate runs inside the loop and its gradients are checked. It is faster than TGLF in the matched linear-solver run and slower than QLKNN under Newton; the QLKNN Newton run itself has stalled steps, so its runtime is not evidence of an equivalent successful solve.
Technical details → Surrogates: in the loop. Decisions → OQ-1, OQ-10, OQ-20.
Technical details → Surrogates › in the loop
Decisions → OQ-1 · OQ-10 · OQ-20