TORAX: A Fast and Differentiable Tokamak Transport Simulator in JAX

We introduce TORAX, an open-source differentiable tokamak core transport simulator targeting fast and accurate core-transport simulation for pulse planning and optimization, and unlocking broad capabilities for controller design and advanced surrogate physics. TORAX is written in Python using JAX, and solves coupled time-dependent 1D PDEs for core ion and electron heat transport, particle transport, and current diffusion. JAX's just-in-time compilation provides fast computation, while maintaining Python's ease of use and extensibility. JAX auto-differentiability enables gradient-based optimization techniques and trajectory sensitivity analysis for controller design, without time-consuming manual Jacobian calculations. JAX's inherent support for neural network development and inference facilitates coupling ML-surrogates of physics models, key for fast and accurate simulation. Code verification is obtained by comparison with the established RAPTOR code on ITER-like and SPARC scenarios. Application of TORAX is demonstrated within a pulse planning workflow designed for optimization of SPARC actuator trajectories. TORAX is an open source tool, and aims to be a foundational component of wider workflows built by the wider community for future tokamak integrated simulations.