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 a code directory where you will install the virtual env and other TORAX dependencies.

mkdir /path/to/torax_dir && cd "$_"

Where /path/to/torax_dir should be replaced by a path of your choice.

Create a TORAX virtual env:

python3 -m venv toraxvenv

Activate the virtual env:

source toraxvenv/bin/activate

It is convenient to set up an alias for the above command.

Install TORAX from PyPI

This is the simplest way to install TORAX. If you don’t plan to do any development work, this is the recommended method.

pip install torax

Install TORAX from Github

If you plan to develop TORAX, we recommend installing from source. See Contribution tips for an installation guide.

JAX environment variables

The following may optionally be added to ~/.bashrc and will cause jax to store compiled functions to the filesystem, avoiding recompilation in some cases:

export JAX_COMPILATION_CACHE_DIR=<path of your choice, such as ~/jax_cache>
export JAX_PERSISTENT_CACHE_MIN_ENTRY_SIZE_BYTES=-1
export JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS=0.0

For more information see Using Jax in TORAX.

(Optional) GPU support

Install additional GPU support for JAX if your machine has a GPU: https://jax.readthedocs.io/en/latest/installation.html#supported-platforms

(Optional) Install QLKNN-hyper

TORAX uses the |qlknn_7_11| transport model by default, an upgrade to the QLKNN-hyper-10D (QLKNN10D) neural network surrogate model of QuaLiKiz |qlknn10d|. QLKNN_7_11 is based on QuaLiKiz 2.8.1 which has an improved collision operator compared to the QLKNN10D training set. QLKNN_7_11 training data also includes impurity density gradients as an input feature and has better coverage of the near-LCFS region compared to QLKNN-hyper-10D.

However, certain use-cases may require the use of QLKNN10D, such as for comparative studies with other simulators. To install QLKNN10D, first download the QLKNN dependencies at a location of your choice:

git clone https://gitlab.com/qualikiz-group/qlknn-hyper.git

To use QLKNN10D , you then need to set model_path in the transport section of your TORAX config to the path of the cloned repository. See Simulation input configuration for more details.

(Optional) Install TGLF wrapper

TORAX supports running the TGLF quasilinear turbulent transport model directly in memory via an f2py extension module (tglf2py). If you want to use TORAX with TGLF, you can compile and install this wrapper manually after installing core TORAX.

Prerequisites

Ensure you have a Fortran compiler (such as gfortran), LAPACK development libraries, OpenMP libraries, and required Python build utilities installed:

pip install meson ninja

Isolate the Fortran Sources

First, ensure you have access to the GACODE repository (available through General Atomics at https://gacode.io). After obtaining or cloning GACODE, export the GACODE_ROOT environment variable to point to your repository root:

export GACODE_ROOT=/path/to/your/gacode

Next, run the following command to isolate the required TGLF Fortran sources while excluding MPI, driver, and unused variants (noting that GACODE filenames like tglf_run.F90 use uppercase extensions):

TGLF_SRCS=$(find "$GACODE_ROOT/tglf/src" -maxdepth 1 -iname "*.f90" \
  | grep -vE \
  -e '/tglf\.f90$' \
  -e '/tglf_mpi\.f90$' \
  -e '/tglf_init_mpi\.f90$' \
  -e '/tglf_driver\.f90$' \
  -e '/tglf_driver_mpi\.f90$' \
  -e '/tglf_TM_mpi\.f90$' \
  -e '/tglf_TM_driver\.f90$' \
  -e '/tglf_nn_TM\.fann\.f90$' \
  -e '\.f90_save$' \
  -e '\.cray\.f90$' \
  -e '\.ieee\.f90$')

Compile the Extension Module

Navigate directly to your installed package directory inside torax/_src/transport_model/tglf and run f2py. Modern f2py will parse the module name straight from your .pyf signature file and invoke Meson under the hood:

TORAX_DIR=$(python -c \
  "import torax, os; print(os.path.dirname(torax.__file__))")
cd "$TORAX_DIR/_src/transport_model/tglf"

FLAGS="-fdefault-real-8 -fdefault-double-8 -ffree-line-length-256 "
FLAGS+="-cpp -fPIC -fopenmp -Ofast -std=f2018 -fall-intrinsics -Wall -W "
FLAGS+="-fimplicit-none -fmax-stack-var-size=65536 -frecursive "
FLAGS+="-falign-commons"

python -m numpy.f2py -c tglf2py.pyf \
  $TGLF_SRCS \
  --f90flags="$FLAGS" \
  --dep lapack \
  -lgomp

Once completed, a shared library binary (e.g., tglf2py_lib*.so or .pyd) will be generated directly inside torax/_src/transport_model/tglf/, which will be used when tglf_run() is called.

Verify the Installation

You can check if the wrapper installed correctly by executing its unit test directly with Python:

python torax/torax/_src/transport_model/tglf/tests/tglf2py_test.py