What is TORAX by Google DeepMind for nuclear fusion?

Meet TORAX, Google DeepMind’s Breakthrough in Open-Source Nuclear Fusion Simulation

TORAX incorporates modular physics-based and machine-learning models, leveraging JAX for fast runtimes with JIT compilation and automatic differentiation, enabling gradient-based optimisation workflows, Jacobian-based PDE solvers, and coupling to machine-learning surrogate models of physics.

Gopika Raj

Published 2025-12-24T09:52:19.010+05:30Updated 2024-06-13T05:30:00.000+05:30

Key takeaways

Launch TORAX, an open-source tokamak core transport simulator, to enhance fusion energy research.

Utilize JAX for fast runtimes and advanced machine learning integration in nuclear fusion simulations.

Achieve high accuracy with TORAX, verified against RAPTOR, showing 1% deviation in plasma profiles.

Leverage gradient-based optimization and machine-learning models to improve tokamak modeling efficiency.

Encourage collaboration in fusion research through open-source access to TORAX code and resources.

Google DeepMind researchers have released TORAX, a new open-source differentiable tokamak core transport simulator implemented in Python using the JAX framework. TORAX essentially simulates the transport of particles, energy, and momentum within the core of a tokamak fusion reactor.

According to the new paper, TORAX solves coupled equations for ion heat, electron heat, particle transport, and current diffusion. It incorporates modular physics-based and machine-learning models, leverages JAX for fast runtimes via just-in-time compilation and automatic differentiation, enables gradient-based optimisation workflows and Jacobian-based PDE solvers, and facilitates coupling to machine-learning surrogate models of physics.

TORAX has been verified against the established RAPTOR code, demonstrating excellent agreement in simulated plasma profiles at stationary state. For an ITER L-mode scenario, the normalised root-mean-square deviation between TORAX and RAPTOR temperature and density profiles was around 1%.

A key innovation is TORAX's use of the JAX framework, allowing just-in-time compilation for speed and automatic differentiation for advanced algorithms like gradient-based optimisation. JAX also simplifies the integration of machine learning surrogate models like the QLKNN neural network trained on gyrokinetic turbulence simulations.

"TORAX offers a powerful and versatile tool for accelerating fusion energy research," said Google DeepMind research scientist and lead author Jonathan Citrin. "Its differentiability and ability to leverage machine learning models are game-changers."

The open-source TORAX code aims to foster collaboration and rapid progress in tokamak modelling for fusion reactor design and operation.

Simulation Training

Google DeepMind has a history of open-sourcing simulators for this purpose. Back in 2020, they released a scalable environment simulator for artificial intelligence research, which helped DeepMind create 2D environments for AI and machine learning research. Simulated training is also the most commonly adopted technique to equip general-purpose robots for the real world.

Meet TORAX, Google DeepMind’s Breakthrough in Open-Source Nuclear Fusion Simulation

Key takeaways

Stay ahead with AI insights.