google-deepmind/torax

google-deepmind/torax

Description: TORAX: Tokamak transport simulation in JAX

Language: Python

License: NOASSERTION

Stars: 679

Forks: 131

Open issues: 65

Created: 2024-03-05T14:20:02Z

Pushed: 2026-06-10T18:19:44Z

Default branch: main

Fork: no

Archived: no

README: 

What is TORAX?

TORAX is a 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 using 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

For more comprehensive documentation, see our readthedocs page.

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

• Finite-volume-method discretization
• Multiple solver options: linear with Pereverzev-Corrigan terms and the

predictor-corrector method, nonlinear with Newton-Raphson, nonlinear 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]](https://meetings.aps.org/Meeting/DPP24/Session/NP12.106) (as-yet covering limited regimes),

• Neoclassical bootstrap current and conductivity with the analytical Sauter

model

• Coupling to the [[QLKNN_7_11]](https://github.com/google-deepmind/fusion_surrogates/)

and QLKNN10D [[van de Plassche et al, Phys. Plasmas 2020]](https://doi.org/10.1063/1.5134126) QuaLiKiz neural network surrogates for physics-based turbulent transport

• General geometry, provided via CHEASE, FBT, or EQDSK equilibrium files
• For testing and demonstration purposes, a single CHEASE equilibrium file

is available in the data/third_party/geo directory. It corresponds to an ITER hybrid scenario equilibrium based on simulations in [[Citrin et al, Nucl. Fusion 2010]](https://doi.org/10.1088/0029-5515/50/11/115007), and was obtained from PINT. A PINT license file is available in data/third_party/geo

• Time dependent geometry is supported by providing a time series of geometry files
• Simple pedestal models using a local adaptive source to set internal boundary

conditions

• Sawtooth triggering and profile redistribution

Additional heating and current drive sources can be provided by user-provided analytical models, or user-provided prescribed data.

Model implementation was verified through direct comparison of simulation outputs to the RAPTOR [[Felici et al, Plasma Phys. Control. Fusion 2012]](https://iopscience.iop.org/article/10.1088/0741-3335/54/2/025002) tokamak transport simulator.

This is not an officially supported Google product.

Development roadmap

A development roadmap is outlined in our readthedocs pages.

Installation guide

Requirements

Install Python 3.11 or greater.

Make sure that tkinter is installed:

sudo apt-get install python3-tk

How to install

Prepare a virtual environment

Install Virtualenv (if not already installed):

pip install --upgrade pip

pip install virtualenv

Create and activate a virtual environment

python3 -m venv toraxvenv
source toraxvenv/bin/activate

Install from PyPI

The simplest way to use TORAX is to install it via PyPI:

pip install torax

You can check that everything runs as it should:

run_torax --config=examples/basic_config.py --quit

If you plan to help develop TORAX, you will need to clone the repository, see contributing and contribution_tips

Running an example

The following command will run TORAX using the default configuration file examples/basic_config.py.

run_torax --config='examples/basic_config.py'

Simulation progress is shown by a terminal progress bar indicating the current time and percentage completed.

To run more involved, ITER-inspired simulations, run:

run_torax --config='examples/iterhybrid_rampup.py'

and

run_torax --config='examples/iterhybrid_predictor_corrector.py'

Additional configuration is provided through flags which append the above run command, and environment variables. For example, for increased output verbosity, you can run with the --log_progress flag.

run_torax --config='examples/iterhybrid_rampup.py' --log_progress

Set environment variables

##### Error checking

If true, error checking is enabled in internal routines. Used for debugging. Default is false since it is incompatible with the persistent compilation cache.

$ export TORAX_ERRORS_ENABLED=

##### JAX Compilation and Cache

If false, JAX does not compile internal TORAX functions. Used for debugging. Default is False.

$ export JAX_DISABLE_JIT=

The following implements the JAX persistent cache and will cause jax to store compiled programs to the filesystem, reducing recompilation time in some cases:

$ export JAX_COMPILATION_CACHE_DIR=
$ export JAX_PERSISTENT_CACHE_MIN_ENTRY_SIZE_BYTES=-1
$ export JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS=0.0

Set flags

log_progress - output simulation time, dt, and number…