Installation Guide

Requirements

Install Python 3.10 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.

Download and install the TORAX codebase via http:

git clone https://github.com/google-deepmind/torax.git

or ssh (ensure that you have the appropriate SSH key uploaded to github).

git clone git@github.com:google-deepmind/torax.git

Enter the TORAX directory and pip install the dependencies.

cd torax; pip install -e .[dev]

TORAX uses the QLKNN_7_11 transport model by default. It can be overridden by specifying a QLKNN model path through the TORAX_QLKNN_MODEL_PATH environment variable. To use the default transport model (recommended), keep the TORAX_QLKNN_MODEL_PATH environment variable empty. Previous versions of TORAX required the environment variable to be set. If you set this variable in a previous TORAX installation, make sure you do not define it in your ~/.bashrc. You can check if the variable is defined by running:

echo $TORAX_QLKNN_MODEL_PATH

If the variable is defined, you can clear it by running:

unset TORAX_QLKNN_MODEL_PATH

For an alternative transport model, see (Optional) Install QLKNN-hyper.

The same applies to TORAX_GEOMETRY_DIR. Previous versions of TORAX required it to be set. It is now recommended to leave the variable empty if using the default geometry. Make sure the definition is removed from ~/.bashrc if you set it in a previous installation.

The following may optionally be added to ~/.bashrc and will cause jax to store compiled programs 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 the Jax persistent cache.

(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

An alternative to QLKNN_7_11 is to use QLKNN-hyper-10D, also known as QLKNN10D (K.L. van de Plassche PoP 2020). 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 includes impurity density gradients as an input feature and has better coverage of the near-LCFS region compared to QLKNN-hyper-10D. However, it is still widely used in other simulators, so it can be useful for comparative studies for instance.

Download QLKNN dependencies:

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

To use this transport model, you need to set the environment variable TORAX_QLKNN_MODEL_PATH to the path of the cloned repository.

export TORAX_QLKNN_MODEL_PATH="$PWD"/qlknn-hyper