TORAX Overview

TORAX is an auto-differentiable tokamak core transport simulator aimed for fast and accurate forward modelling, pulse-design, trajectory optimization, and controller design workflows. TORAX is written in Python-JAX, with the following motivations:

  • Open-source and extensible, aiding with flexible workflow coupling.

  • JAX provides auto-differentiation capabilities and code compilation for fast runtimes. Differentiability allows for gradient-based nonlinear PDE solvers for fast and accurate modelling, and for sensitivity analysis of simulation results to arbitrary parameter inputs, enabling applications such as trajectory optimization and data-driven parameter identification for semi-empirical models. Auto-differentiability allows for these applications to be easily extended with the addition of new physics models including ML-surrogates, or new parameter inputs, by avoiding the need to hand-derive Jacobians

  • Python-JAX is a natural framework for the coupling of ML-surrogates of physics models.

TORAX, at v1.0.0, has the following physics and numerics feature set:

  • Coupled PDEs of ion and electron heat transport, electron particle transport, and current diffusion, solved numerically with:

    • Finite-volume-method spatial discretization

    • Multiple solver options for PDE time evolution

      • Linear with Pereverzev-Corrigan terms, and predictor-corrector steps for approximating the nonlinear solution.

      • Nonlinear solver with Newton-Raphson iterations, with line search and timestep backtracking.

      • Nonlinear solver with optimization using the jaxopt library.

      • Poloidal flux boundary conditions based on either total current or loop voltage at the last-closed-flux-surface.

  • Ohmic power, ion-electron heat exchange, fusion power, Bremsstrahlung, impurity line radiation, an ICRH ML-surrogate [Wallace et al, APS 2024] (as-yet covering limited regimes).

  • Bootstrap current and neoclassical conductivity with the analytical Sauter model.

  • Coupling to the QLKNN_7_11 and QLKNN10D [van de Plassche et al, Phys. Plasmas 2020] QuaLiKiz-neural-network surrogates for physics-based turbulent transport. The semi-empirical Bohm-GyroBohm model is also available.

  • General geometry, provided via CHEASE, FBT, or EQDSK equilibrium files.

  • Sawtooth triggering and profile redistribution.

  • Simple pedestal models using a local adaptive source to set internal boundary conditions.

  • API for setting time-dependent boundary conditions, prescribed sources, and prescribed profiles optionally evolved by the PDE, from a variety of input data structures or user-provided analytical models.

    • For testing and demonstration purposes, a single CHEASE equilibrium file is available in the data/geo directory. It corresponds to an ITER hybrid scenario equilibrium based on simulations in [Citrin et al, Nucl. Fusion 2010] and was obtained from PINT. A PINT license file is available in data/third_party/geo.

Model implementation was verified through direct comparison of simulation outputs to the RAPTOR [Felici et al, Plasma Phys. Control. Fusion 2012] tokamak transport simulator.