A Differentiable Tokamak Simulator That Runs on JAX

If you've ever worked on fusion plasma modeling, you know the drill: you've got a set of coupled partial differential equations, you're hand-deriving Jacobians, and every time you want to test a new physics model, you're back to writing gradients by hand. It's slow, error-prone, and doesn't play nicely with machine learning workflows. TORAX, a project from Google DeepMind, takes a different approach—it builds a tokamak core transport simulator on top of JAX, which means gradients come for free.

What It Does

TORAX is a differentiable tokamak core transport simulator written in Python using JAX. It's designed for forward modeling, pulse design, trajectory optimization, and controller design workflows. At version 1.0.0, it solves coupled PDEs for ion and electron heat transport, electron particle transport, and current diffusion.

The numerics side includes finite-volume-method discretization and several solver options: linear solvers with Pereverzev-Corrigan terms and the predictor-corrector method, plus nonlinear solvers using Newton-Raphson or optimization via the jaxopt library. Poloidal flux boundary conditions can be set based on either total current or loop voltage at the last-closed-flux-surface.

On the physics side, you get Ohmic power, ion-electron heat exchange, fusion power, Bremsstrahlung, impurity line radiation, and an ICRH ML-surrogate (though the README notes it covers limited regimes). Neoclassical bootstrap current and conductivity use the analytical Sauter model. Turbulent transport can be handled through coupling to QLKNN_7_11 and QLKNN10D QuaLiKiz neural network surrogates. Geometry is provided via CHEASE, FBT, or EQDSK equilibrium files, and time-dependent geometry is supported by supplying a time series of geometry files. There are also simple pedestal models and sawtooth triggering with profile redistribution.

Why It's Cool

• Differentiability isn't a gimmick here. Because TORAX is built on JAX, you get auto-differentiation and code compilation for fast runtimes. That means gradient-based nonlinear PDE solvers, sensitivity analysis of simulation results to arbitrary parameter inputs, and the ability to extend these workflows to new physics models without hand-deriving Jacobians. If you've ever spent a week deriving derivatives for a new model, you'll appreciate what that saves.

• ML-surrogates fit naturally. The Python-JAX stack is a natural framework for coupling machine learning surrogates of physics models. The project already couples to QuaLiKiz neural network surrogates for turbulent transport, and the architecture makes it straightforward to add more.

• Multiple solver options give you flexibility. Rather than forcing one numerical approach, TORAX offers linear and nonlinear solvers, so you can pick what works for your problem. The boundary conditions are similarly flexible—total current or loop voltage, your choice.

• It's verified against RAPTOR. Model implementation was verified through direct comparison of simulation outputs to the RAPTOR tokamak transport simulator. That's a meaningful validation step, not just a claim.

• Open-source and extensible. The project is designed to aid flexible workflow coupling. You can provide additional heating and current drive sources through user-provided analytical models or prescribed data. And if you're testing or demonstrating, there's a single CHEASE equilibrium file in the data/third_party/geo directory corresponding to an ITER hybrid scenario equilibrium.

• Time-dependent geometry. Providing a time series of geometry files lets you model evolving plasma shapes, which matters for realistic pulse design.

How to Try It

You'll need Python 3.12 or greater and tkinter installed. On Debian-based systems:

sudo apt-get install python3-tk

The README recommends setting up a virtual environment (the installation guide continues from there—check the full README for the complete steps). Once installed, you can start exploring the simulator with the provided CHEASE equilibrium file and the documented API.

For comprehensive documentation, including the development roadmap, head to the readthedocs page.

The repository is at https://github.com/google-deepmind/torax. Note that this is not an officially supported Google product.

Final Thoughts

TORAX is a solid tool for researchers and engineers working on tokamak transport who want to incorporate gradient-based optimization or machine learning into their workflows. The JAX foundation makes it a natural fit if you're already in that ecosystem, and the solver and physics options give you room to tackle different problems. It's not a turnkey solution—you'll need domain knowledge to use it well—but if you're doing pulse design, trajectory optimization, or controller design for tokamaks, the differentiability alone is worth a look. The project is at v1.0.0, so it's mature enough to build on, and the extensibility means you can grow it with your own physics models.

Follow @githubprojects for more developer tools and open source projects.