🔢 mathematics

TokaGrad: End-to-end differentiable tokamak simulator for L-to-H full scenario optimization

This paper introduces TokaGrad, the first end-to-end differentiable tokamak simulator that self-consistently models full-discharge scenarios—including ramp-up, L-mode, and H-mode transitions—to enable efficient, gradient-based optimization of reactor designs and autonomous plasma control.

Original authors: Jaemin Seo

Published 2026-07-13

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jaemin Seo

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine trying to bake the perfect cake, but instead of a recipe, you have to guess the temperature, the mixing speed, and the ingredients by throwing darts at a board. If the cake burns, you try again. If it's too dry, you try again. This is basically how scientists have traditionally designed fusion reactors—massive, doughnut-shaped machines called tokamaks that aim to replicate the Sun's power on Earth. They rely on expensive "trial and error," tweaking knobs and hoping for the best.

Enter TokaGrad, a new digital tool developed by researcher Jaemin Seo. Think of TokaGrad not as a black box that you poke blindly, but as a super-smart, self-aware video game engine for fusion.

The Magic of the "Self-Aware" Simulator

In the old days, simulating a fusion reaction was like watching a movie where you couldn't rewind or pause to see why a character made a specific move. You just saw the result. If the plasma (the super-hot gas inside the reactor) didn't behave, scientists had to guess what went wrong.

TokaGrad changes the rules. It builds the entire simulation as a connected chain of dominoes. Every single part of the machine—from the shape of the magnetic fields to the heating power—is linked. If you nudge one domino (like changing the heating power), the tool can instantly calculate exactly how that tiny nudge ripples all the way through the system to change the final result.

This is called differentiable programming. Instead of guessing, the tool knows the "slope" of the problem. It's like having a GPS that doesn't just tell you where you are, but instantly calculates the exact steering angle needed to get to your destination, no matter how twisty the road gets.

The Big Challenge: The "L-to-H" Jump

Fusion reactors have a tricky personality. They can operate in a lazy mode (called L-mode) or a high-performance mode (called H-mode). The H-mode is the goal because it traps heat much better, but getting there is like trying to push a heavy boulder over a hill. You need just the right amount of push at just the right time.

If you push too hard too soon, you might not get over the hill. If you push too little, you slide back. In the real world, finding this perfect push is incredibly hard because the physics are messy and non-linear.

TokaGrad is the first simulator that can model this entire journey—from the cold start-up, through the lazy L-mode, all the way to the high-performance H-mode—while keeping those "domino links" intact. It doesn't just simulate the final state; it simulates the transition.

What the Tool Actually Did (The Proof)

The researchers didn't just build the tool; they tested it in three exciting ways:

The Speed Test: They compared TokaGrad to an existing tool called TORAX. While TORAX is like a slow, careful accountant, TokaGrad is a sprinter. On a standard laptop computer, TokaGrad ran a simulation more than ten times faster than TORAX. It was so fast that it could simulate a 150-second fusion pulse in just 17 seconds of real time.
The "Smart" Control: They asked the tool to figure out how to get the most energy out of a reactor. Instead of trying random settings, the tool used its "slope" knowledge to adjust the heating and current. It found a solution where it increased the heating power temporarily to jump the H-mode hill, then dialed it back down to save energy while keeping the high performance. It did this automatically, without human guesswork.
The "What-If" Designer: They let the tool redesign the reactor itself. By tweaking the size and shape of the machine (like the major radius and magnetic field strength), the tool suggested a new design that could produce even more energy. It found that making the reactor slightly larger and changing its shape could boost performance, all while keeping the machine stable.

What This Means (And What It Doesn't)

The paper suggests that this approach could revolutionize how we design fusion reactors. Instead of spending years on trial-and-error experiments, engineers could use TokaGrad to instantly test thousands of scenarios and find the best path to a working reactor.

However, it's important to remember that these are simulations. The tool has shown it works on a computer, predicting that it can find better heating strategies and reactor shapes. It has not yet built a physical reactor or run a real fusion experiment. The "breakthrough" here is the method of optimization, not a physical fusion reactor sitting in a lab today.

The authors also note that while the tool is fast and smart, it still relies on the physics models fed into it. If the models are wrong, the "smart" advice might be wrong too. But for now, TokaGrad offers a promising new way to navigate the complex, bumpy road toward clean, limitless fusion energy. It turns the process of designing a star on Earth from a game of chance into a game of precise, calculated strategy.

Technical Summary: TokaGrad – End-to-End Differentiable Tokamak Simulator

1. Problem Statement

As magnetic fusion research transitions from scientific feasibility to reactor-scale commercialization, the design of new reactor concepts and the optimization of high-performance plasma scenarios have become critical. Traditional approaches to scenario development rely heavily on costly trial-and-error, expert-driven waveform tuning, and low-dimensional parameter scans. While integrated modeling tools (e.g., ASTRA, TRANSP) have improved virtual experimentation, scenario design remains a "human-in-the-loop" process.

The primary limitations of current methods are:

Black-box Optimization: Many machine learning approaches (Bayesian optimization, reinforcement learning) treat the simulator as a black box, requiring repeated forward evaluations and indirect sensitivity inference, which becomes computationally prohibitive for high-dimensional, multi-actuator, full-discharge scenarios.
Lack of Differentiability: Conventional numerical simulations are executed as step-by-step procedures where gradients cannot propagate through the entire pipeline. This prevents direct gradient-based control and optimization using internal sensitivities.
Incomplete Full-Scenario Modeling: Existing differentiable transport solvers (e.g., TORAX) have primarily focused on core transport and prescribed scenarios. They lack self-consistent modeling of dynamic full-discharge scenarios involving L-to-H transitions, edge pedestal formation, and the coupling of evolving plasma shapes with equilibrium reconstruction.

2. Methodology

The authors introduce TokaGrad, an end-to-end differentiable tokamak transport simulator built on the JAX framework. The system represents the entire simulation pipeline—from initial conditions to final plasma performance—as a connected differentiable computational graph. This allows automatic differentiation (AD) to compute Jacobians across the full simulation, linking machine parameters, actuator waveforms, and plasma responses.

Core Architecture

TokaGrad integrates the following differentiable modules within a single graph:

Transport Framework: The core consists of coupled one-dimensional transport equations for electron/ion temperature ( $T_{e,i}$ ), electron density ( $n_e$ ), and poloidal magnetic flux ( $\psi$ ) in a normalized radial flux coordinate ( $\rho$ ).
- Numerical Scheme: A semi-implicit finite-volume method is used for stability, solved via a differentiable tridiagonal matrix solver (Thomas algorithm) compatible with JAX's just-in-time (JIT) compilation.
- Equations: Includes particle transport, conservative energy transport (for $T_{e,i}$ ), and inductive flux evolution coupled to bootstrap and current-drive sources.
Geometry and Equilibrium:
- EQDSK Mode: Supports externally generated equilibria (e.g., from EFIT) for realistic benchmarking, though gradients with respect to geometry are limited unless parameterized.
- Miller/Moment Mode: A reduced fixed-boundary model using analytic flux-surface parameterization (major radius, elongation, triangularity, Shafranov shift). This mode is fully differentiable, enabling optimization of machine shape parameters ( $R_0, a, \kappa$ ).
Transport Closures:
- Turbulent Transport: Uses Neural Network (NN) surrogates (TGLF-NN, QLKNN) to replace computationally expensive gyrokinetic solvers. These surrogates provide differentiable mappings from plasma profiles to transport coefficients.
- Neoclassical Transport: Supports analytic formulas (Sauter-Angioni) and NN surrogates (NEO-NN).
- Empirical Models: Includes a Bohm-gyroBohm (BgB) model for rapid baselines.
Sources and Sinks:
- Heating: Ohmic, alpha heating (using Bosch-Hale fit), collisional exchange, and radiation losses. Auxiliary heating is modeled via smooth, normalized Gaussian profiles.
- Particle Sources: Implements density feedback control (tracking Greenwald fraction) or prescribed-shape models to avoid non-differentiable hard switching.
L-H Transition and Pedestal:
- Transition: Uses a smooth sigmoid function to gate the transition between L-mode and H-mode based on the power crossing the separatrix ( $P_{sep}$ ) relative to an empirical threshold ( $P_{LH}$ , e.g., Martin or Delabie scalings). This avoids non-differentiable binary switches.
- Pedestal: Constructs pedestal targets using either an alpha-critical reduced model or an EPED1-NN surrogate. The pedestal profile is blended smoothly with the core transport solution using a smoothmax function, ensuring gradients propagate through the transition.

3. Key Contributions

First Self-Consistent Differentiable Full-Discharge Simulator: TokaGrad is the first framework capable of modeling dynamic full-discharge scenarios where actuators, plasma equilibrium, pedestal formation, and confinement-regime transitions (L-to-H) evolve together within a differentiable graph.
End-to-End Differentiability: By converting high-fidelity physics closures (TGLF, EPED1, NEO) into JAX-compatible NN surrogates, the system enables gradients to propagate from final performance metrics back through transport, turbulence, and equilibrium models to design and control inputs.
Interactive Live Simulation: The framework supports faster-than-real-time simulation (demonstrated as >10x speedup over TORAX for long pulses after JIT compilation), enabling interactive exploration of cause-and-effect relationships.
Gradient-Based Optimization: The system facilitates direct optimization of actuator waveforms, scenario trajectories, and reactor design parameters without requiring finite-difference scans.

4. Results and Demonstrations

The paper validates TokaGrad through benchmarks and optimization case studies:

Benchmarking:
- Against TORAX: In a simplified L-mode scenario with circular geometry, TokaGrad reproduced core transport profiles (temperature, density, current) within ~5% error of the established JAX-based TORAX simulator.
- Performance: TokaGrad demonstrated a wall-clock time speedup of more than one order of magnitude compared to TORAX for long-pulse simulations once JIT compilation overhead was amortized.
- H-Mode Simulation: Simulations of the ITER baseline scenario using TGLF-NN/EPED1-NN and QLKNN/alpha-critical models achieved fusion gain $Q \sim 11$ and fusion power $\sim 550$ MW in ~7 seconds (faster than real-time).
Full-Discharge L-to-H Simulation:
- A 150-second ITER baseline simulation successfully modeled the ramp-up, L-mode phase, and transition to H-mode.
- The simulation captured the hysteresis of the L-H transition: a transient heating overshoot (73 MW) triggered the transition, after which alpha heating sustained H-mode even when auxiliary power was reduced to 50 MW.
Optimization Applications:
- Actuator Control: Gradient-based control (Adam optimizer) successfully adjusted plasma current ( $I_p$ ), auxiliary heating ( $P_{aux}$ ), and Greenwald fraction ( $f_{GW}$ ) to maximize fusion metrics ( $Q + 0.01 P_{fus}$ ), respecting stability constraints ( $q_{95} > 3, \beta_N < 3$ ).
- Scenario Waveform Optimization: Starting from a naive ramp-up that failed to access H-mode, the optimizer automatically discovered a transient heating overshoot strategy. This allowed the plasma to cross the L-H threshold and sustain high-performance H-mode under the same final actuator conditions.
- Reactor Design Optimization: The framework optimized machine design parameters ( $R_0, a, B_T, I_p$ ) to maximize fusion performance. The optimizer correctly identified trends such as increasing major radius and magnetic field, while adjusting parameters to satisfy stability constraints.

5. Significance and Claims

The authors claim that TokaGrad provides a practical route toward automated tokamak scenario design. By replacing black-box optimization with gradient-based methods enabled by differentiable physics, the framework allows for:

Direct Sensitivity Analysis: Computing how fusion gain depends on heating trajectories or how pedestal height depends on geometry without external perturbation.
Path-Dependent Optimization: Optimizing not just the final state, but the entire dynamical pathway (e.g., ramp-up waveforms) required to reach high-performance regimes like H-mode.
Scalability: The ability to perform repeated forward and backward evaluations on laptop-scale hardware (Apple M1 Pro) makes it feasible to optimize complex, multi-objective reactor scenarios that were previously computationally impractical.

The paper concludes that while current demonstrations focus on ITER-baseline-like conditions, the framework is general and offers a new pathway for automated, physics-informed optimization of burning-plasma scenarios and next-generation reactor concepts.