🔬 physics

TORAX: A Fast and Differentiable Tokamak Transport Simulator in JAX

This paper introduces TORAX, an open-source, differentiable tokamak transport simulator built in JAX that leverages automatic differentiation and JIT compilation to enable fast, gradient-based optimization for plasma scenario modeling, pulse design, and control.

Original authors: Jonathan Citrin, Ian Goodfellow, Akhil Raju, Jeremy Chen, Jonas Degrave, Craig Donner, Federico Felici, Philippe Hamel, Andrea Huber, Dmitry Nikulin, David Pfau, Brendan Tracey, Martin Riedmiller, Pus

Published 2026-07-10

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Jonathan Citrin, Ian Goodfellow, Akhil Raju, Jeremy Chen, Jonas Degrave, Craig Donner, Federico Felici, Philippe Hamel, Andrea Huber, Dmitry Nikulin, David Pfau, Brendan Tracey, Martin Riedmiller, Pushmeet Kohli

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 you're trying to bake the perfect loaf of bread, but instead of flour and water, your ingredients are super-hot plasma, magnetic fields, and fusion energy. The goal? To build a machine (a tokamak) that can power the world. The problem is, predicting how that plasma behaves inside the machine is like trying to forecast the weather in a hurricane while the storm is happening. It's messy, it's fast, and it involves a dizzying mix of physics rules that change every millisecond.

Enter TORAX, a new, open-source computer program created by researchers at Google DeepMind. Think of TORAX as a "smart, super-fast flight simulator" for these fusion reactors. But here's the twist: unlike old simulators that are like clunky, slow-moving trucks, TORAX is built with a special engine called JAX. This engine allows the simulator to not only run incredibly fast but also to "learn" from its mistakes instantly by calculating gradients (basically, figuring out exactly which knob to turn to get a better result).

The Old Way vs. The New Way

For years, scientists have used other simulators (like RAPTOR, JETTO, or ASTRA) to model these reactors. These tools are like reliable, heavy-duty calculators. They work well, but they are often written in older languages (like Fortran) that make it hard to plug in modern Artificial Intelligence (AI) models or to quickly tweak the math to find the perfect settings. They are also slow; running a simulation for a few seconds of reactor time could take hours or even days on a standard computer.

The authors of this paper argue that this slowness and lack of flexibility are holding fusion research back. They explicitly rule out the idea that we should stick with these old, rigid codes if we want to use modern machine learning or do complex optimization. They show that while old codes are good, they can't easily handle the "what-if" scenarios needed to design future reactors or control them in real-time.

How TORAX Works: The Magic of the "Flight Simulator"

TORAX simulates the core of the plasma by solving a set of four main equations. You can think of these as the rules governing:

Ion Heat: How hot the heavy particles get.
Electron Heat: How hot the light particles get.
Particle Transport: How the density of the plasma changes.
Current Diffusion: How the electric current flows through the plasma.

Instead of solving these one by one in a slow, step-by-step march, TORAX uses a technique called Finite Volume Method. Imagine slicing the plasma into a stack of thin, invisible pancakes (or cells). The program calculates how heat and particles flow from one pancake to the next. It uses a clever math trick called the "power-law scheme" to decide whether to look at the flow from the left or the right, depending on how fast the plasma is moving. This ensures the simulation doesn't get "wobbly" or inaccurate.

The real magic happens because TORAX is written in Python using JAX. This means:

Speed: It compiles the code into a super-efficient version that runs on CPUs, GPUs, or even TPUs. In tests, it ran a simulation of an 80-second reactor startup in just 22 seconds on a standard computer. That's faster than real-time!
Differentiability: This is the big deal. TORAX can calculate exactly how a tiny change in the input (like turning up the heating power by 1%) affects the final result. This is crucial for using AI to automatically design the perfect "pulse" (the sequence of events) for a reactor.

The "Surrogate" Trick: AI as a Shortcut

One of the most exciting features is how TORAX handles turbulence. Turbulence in plasma is incredibly complex, like trying to predict every swirl in a river. Usually, simulating this takes forever. TORAX can use Machine Learning (ML) surrogates instead.

Think of an ML surrogate as a "crystal ball" trained on millions of complex physics simulations. Instead of doing the heavy lifting of calculating every swirl, the crystal ball (the AI model) instantly predicts the result. TORAX is built to plug these AI models in seamlessly. For example, it uses a model called QLKNN, which is trained on data from a more complex code called QuaLiKiz. This allows TORAX to run simulations with high-fidelity turbulence models in a fraction of the time it used to take.

Did It Work? The Proof is in the Pudding

The authors didn't just build it; they tested it. They ran TORAX side-by-side with RAPTOR, a well-established and trusted simulator. They simulated an "ITER-like" scenario (ITER is the massive international fusion experiment being built in France).

The results were impressive:

Agreement: The temperature, density, and magnetic profiles generated by TORAX matched RAPTOR almost perfectly. The difference was tiny—less than 1% for steady states and around 2.5% during fast changes.
Speed: While RAPTOR is fast, TORAX was able to run the same scenarios even faster, especially when using the AI surrogates.
Verification: The paper explicitly states that TORAX is "verified" against RAPTOR, meaning the results are trustworthy within the tested scenarios.

What's Next? (And What's Not There Yet)

The paper is careful not to overhype. It admits that TORAX is a "new" tool and that some things are still on the roadmap.

What's missing: It doesn't yet handle complex 3D shapes of the magnetic field dynamically (it uses pre-calculated shapes), nor does it include all types of radiation losses or heavy impurity transport yet.
What's coming: The authors plan to add more physics models, like those for "sawteeth" (a type of magnetic instability) and better ways to handle the edge of the plasma (the "pedestal").

The Bottom Line

TORAX isn't a magic wand that solves fusion overnight. It's a powerful, flexible, and incredibly fast tool that makes the job of designing fusion reactors much easier. By combining the speed of modern computing with the smarts of machine learning, it allows scientists to run thousands of "what-if" scenarios in the time it used to take to run just one. It's like upgrading from a hand-cranked calculator to a supercomputer that can also teach itself how to solve the problem better.

The paper concludes that this tool is ready to help accelerate research in tokamak scenario modeling, pulse design, and control, but it emphasizes that it is a step forward in the journey, not the final destination. The future of fusion is bright, and tools like TORAX are helping to light the way.

Technical Summary: TORAX – A Fast and Differentiable Tokamak Transport Simulator in JAX

Problem Statement

Accurate integrated plasma simulation is critical for the design and operation of toroidal magnetic confinement devices like tokamaks, serving applications ranging from experimental interpretation to controller design. However, existing core transport codes (e.g., JETTO, ASTRA, PTRANSP) face a trade-off between physics fidelity and simulation speed; high-fidelity models often require hours to days to simulate a few energy confinement times. Furthermore, traditional codes, typically written in Fortran or MATLAB, lack native differentiability, hindering gradient-based optimization, sensitivity analysis, and seamless integration with machine learning (ML) surrogates. While some dedicated codes like RAPTOR offer speed, they rely on manually coded Jacobians, limiting extensibility to complex ML models. There is a distinct need for a simulator that combines fast runtimes, automatic differentiability, and native support for ML integration.

Methodology

The authors present TORAX, an open-source tokamak core transport simulator implemented in Python using the JAX framework.

Governing Equations and Discretization

TORAX solves a coupled set of 1D transport Partial Differential Equations (PDEs) in normalized toroidal flux coordinates ( $\hat{\rho}$ ). The system governs:

Ion heat transport ( $T_i$ )
Electron heat transport ( $T_e$ )
Electron particle transport ( $n_e$ )
Current diffusion (poloidal flux $\psi$ )

The equations are derived from flux-surface-averaged moments of kinetic equations and Faraday's law. Spatial discretization is performed using a Finite Volume Method (FVM) on a uniform grid. The scheme employs a power-law weighting for convection-diffusion fluxes to smoothly transition between central differencing and upwinding based on the Péclet number.

Solver Implementation

Time integration utilizes the theta method, a weighted average between explicit and implicit Euler schemes. TORAX offers three solver options to handle the resulting nonlinear algebraic systems:

Linear Solver (Predictor-Corrector): Uses fixed-point iteration. To enhance stability with stiff coefficients (e.g., from ML models), it optionally implements the Pereverzev-Corrigan method, adding artificial diffusion balanced by convection.
Newton-Raphson Solver: A gradient-based root-finding method. Crucially, it leverages JAX's automatic differentiation to compute the Jacobian of the residual vector automatically, eliminating the need for manual Jacobian coding.
Optimizer Solver: Recasts the residual as a loss function minimized via iterative optimization algorithms (using jaxopt).

Physics Models

TORAX features a modular framework for physics components:

Geometry: Supports CHEASE equilibrium data and a simplified circular geometry.
Transport: Includes constant, Critical Gradient Model (CGM), and QLKNN, a machine learning surrogate trained on the QuaLiKiz gyrokinetic code. The QLKNN model is natively implemented in JAX using the Flax library.
Sources: Handles ion-electron heat exchange, fusion power (Bosch-Hale parameterization), Ohmic heating, and particle sources (gas puff, pellet injection, NBI). Sources can be treated as implicit or explicit.
Neoclassical Physics: Uses the Sauter model for bootstrap current and neoclassical conductivity.

Software Architecture

The codebase is structured into distinct modules (Geometry, Transport Models, Sources, Steppers, Time Step Calculators, State). It utilizes JAX's Just-In-Time (JIT) compilation to accelerate computationally intensive components (source models, transport coefficients, linear algebra) while keeping control flow in standard Python for flexibility. This design allows for execution on CPUs, GPUs, and TPUs.

Key Results

The authors verified TORAX against the established RAPTOR code, which has itself been validated against other transport solvers.

Stationary State Benchmark (ITER L-mode):
- Simulated 10 seconds of an ITER-like scenario with constant current and heating.
- Results showed excellent agreement between TORAX and RAPTOR for ion/electron temperatures, density, poloidal flux, and safety factor profiles.
- Normalized Root-Mean-Square Deviation (NRMSD) values were approximately 0.7% to 1.1% for stationary profiles.
- Transient temperature deviations were limited to ~2.5%.
Dynamic Scenario Benchmark:
- Simulated a 300-second scenario with time-varying plasma current (ramp-up, flattop, ramp-down) and modulated heating power.
- TORAX and RAPTOR tracked the dynamics closely, with NRMSD values of ~3.5% to 4.0% for temperature profiles over both energy confinement and current diffusion timescales.
Performance:
- On a 48-core CPU node, a simulation of 80 seconds of ITER ramp-up (using QLKNN and Newton-Raphson) completed in 22 seconds of runtime, with a compilation time of 15.6 seconds.
- This demonstrates faster-than-realtime capability for specific configurations, though the authors note this is not a general conclusion for all transient phenomena.

Significance and Claims

The paper positions TORAX as a versatile tool designed to accelerate research in tokamak scenario modeling, pulse design, and control. Its primary significance lies in three areas:

Differentiability: By leveraging JAX, TORAX provides a differentiable PDE residual. This enables gradient-based optimization and sensitivity analysis (via Forward Sensitivity Analysis) without the manual Jacobian coding required by previous differentiable codes like RAPTOR.
ML Integration: The native support for neural networks within the JAX ecosystem facilitates the seamless coupling of ML surrogates (like QLKNN) as physics models, allowing them to be developed in the same language as the simulator.
Speed and Flexibility: JIT compilation ensures fast runtimes, while the Python-based modular architecture simplifies code extensibility, maintenance, and coupling with external workflows.

The authors conclude that TORAX offers a "powerful and versatile tool" for the fusion community, specifically targeting applications that require both speed and differentiability, such as uncertainty quantification, scenario optimization, and model-based controller design. Future work is planned to expand physics fidelity (e.g., MHD, radiation, multi-ion species) and enhance differentiability features (e.g., persistent compilation caches, forward sensitivity calculations).