Skip to main content
Nubint
QUICK REVIEW

[Paper Review] TORAX: A Fast and Differentiable Tokamak Transport Simulator in JAX

J. Citrin, Ian Goodfellow|arXiv (Cornell University)|Jun 10, 2024
Magnetic confinement fusion research6 citations
TL;DR

TORAX is an open-source, differentiable 1D tokamak core transport simulator built in Python with JAX, enabling fast runs and gradient-based optimization, and it benchmarks against RAPTOR.

ABSTRACT

We present TORAX, a new, open-source, differentiable tokamak core transport simulator implemented in Python using the JAX framework. TORAX solves the coupled equations for ion heat transport, electron heat transport, particle transport, and current diffusion, incorporating modular physics-based and ML models. JAX's just-in-time compilation ensures fast runtimes, while its automatic differentiation capability enables gradient-based optimization workflows and simplifies the use of Jacobian-based PDE solvers. Coupling to ML-surrogates of physics models is greatly facilitated by JAX's intrinsic support for neural network development and inference. TORAX is verified against the established RAPTOR code, demonstrating agreement in simulated plasma profiles. TORAX provides a powerful and versatile tool for accelerating research in tokamak scenario modeling, pulse design, and control.

Motivation & Objective

  • Provide a differentiable, fast core tokamak transport simulator implemented in Python/JAX.
  • Enable coupling with ML surrogates and gradient-based optimization workflows.
  • Benchmark TORAX against established codes (e.g., RAPTOR) to verify accuracy and reliability.
  • Offer modular physics models for geometry, transport, sources, and neoclassical physics to support scenario modeling and control.
  • Lay out a roadmap for extensions to include more physics and dynamic equilibria.

Proposed method

  • Formulate a coupled set of 1D transport PDEs for ion/electron heat, electron density, and current diffusion in normalized flux coordinates.
  • Discretize with a finite volume method on a uniform 1D grid and apply a power-law Péclet weighting for face values.
  • Solve time evolution with theta-method based time stepping and provide linear, Newton-Raphson, and optimizer solvers with auto-differentiation via JAX.
  • Incorporate modular physics models including geometry (CHEASE or Circular), transport (Constant, CGM, QLKNN10D), and neoclassical physics (Sauter model).
  • Enable ML-surrogate coupling (QLKNN10D) with dedicated JAX inference code and smoothing options for solver stability.
Figure 1: Poloidal cross section of the core region of a toroidal tokamak plasma, with closed nested magnetic flux surfaces. R is the major radius from the center of the torus, Z is the vertical direction. ITER-like magnitudes are shown. The flux surfaces are elliptical in the poloidal plane for ill
Figure 1: Poloidal cross section of the core region of a toroidal tokamak plasma, with closed nested magnetic flux surfaces. R is the major radius from the center of the torus, Z is the vertical direction. ITER-like magnitudes are shown. The flux surfaces are elliptical in the poloidal plane for ill

Experimental results

Research questions

  • RQ1Can TORAX reproduce core tokamak transport physics with differentiable, fast solvers suitable for gradient-based workflows?
  • RQ2How well does TORAX agree with established transport codes like RAPTOR for plasma profiles?
  • RQ3What is the impact of modular physics models (geometry, transport, neoclassical) on simulation speed, stability, and extensibility?
  • RQ4To what extent can TORAX integrate ML surrogates for transport models while maintaining differentiability and numerical robustness?
  • RQ5What roadmap extensions are planned to broaden TORAX’s physics coverage and scenarios?

Key findings

  • TORAX provides a differentiable, fast core transport simulator in Python using JAX, enabling gradient-based optimization and ML coupling.
  • TORAX is verified against RAPTOR, demonstrating agreement in simulated plasma profiles.
  • TORAX supports multiple physics modules (geometry, transport models CGM/QLKNN10D, neoclassical Sauter) within a modular framework.
  • The solver offers linear, Newton-Raphson, and optimizer-based options with adaptive timestep and line-search/backtracking for robustness.
  • JAX enables automatic differentiation and backend flexibility (CPU/GPU), facilitating ML surrogates and multi-workflow integration.
  • The codebase is open-source with planned extensions to include additional physics and dynamic equilibria.
Figure 2: 1D spatial discretization with the finite volume method, with $N=5$ and $\hat{\rho}=0.2$ . The cell volume centered at $\hat{\rho}=0.5$ is highlighted. Keeping with 3D FVM formalism, a cell volume is defined with $V=d\hat{\rho}A$ , with arbitrary face area $A$ . The corresponding cell face
Figure 2: 1D spatial discretization with the finite volume method, with $N=5$ and $\hat{\rho}=0.2$ . The cell volume centered at $\hat{\rho}=0.5$ is highlighted. Keeping with 3D FVM formalism, a cell volume is defined with $V=d\hat{\rho}A$ , with arbitrary face area $A$ . The corresponding cell face

Better researchstarts right now

From paper design to paper writing, dramatically reduce your research time.

No credit card · Free plan available

This review was created by AI and reviewed by human editors.