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

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.