TORAX code structure

This page describes the structure of the TORAX repository and gives pointers for where to find different components in the code base.

This page is just an overview. For more complete details of what’s available to call and use, see the API docs. And for more information on the logical flow of TORAX, read through the code and reach out to our team if anything is unclear! We strive to make the code readable and understandable with the help of this guide, so feedback is appreciated :)

Runnable entrypoint

run_simulation_main.py is the main runnable entrypoint for trying out TORAX. If you are interested in getting started quickly, check out Quickstart to Running and Plotting.

TORAX library

The TORAX library contains modules to build and run a transport simulation. This section describes several of the modules and how they fit together.

For reference, the standard control flow while using TORAX is:

Configure a simulation with a dictionary of Python primitives, numpy arrays, or xarray datasets (see torax/examples/). More on this config language below.
Use torax.config.build_sim.build_sim_from_config() to convert the config to a torax.sim.Sim object, a helper object containing the building blocks of a TORAX simulation run.
torax.sim.Sim internally calls torax.sim.run_simulation(), which is an entrypoint for running an entire TORAX time loop, Technically, the above steps are optional (though useful standards). The API is designed to enable users to customize TORAX with purpose-built timeloops calling the Stepper. Tutorials illustrating how to customize TORAX timeloops are coming soon.
torax.sim.run_simulation() steps the simulation in a loop, the duration of each time step coming from a TimeStepCalculator, and the state updates coming from a Stepper, which internally uses SourceModels, a TransportModel and a Geometry to help evolve the state.

The rest of this section details each of these components and a few more.

sim.py

sim.py holds the majority of the business logic running TORAX. run_simulation() is its main function which actually runs the TORAX simulation. All TORAX runs using the standard flow hit this function.

This file also contains the torax.sim.Sim class which wraps run_simulation() for convenience. It provides a higher level API so users can run a TORAX simulation without constructing all the inputs run_simulation() asks for.

Broadly, run_simulation() takes the following inputs:

An initial state
A geometry provider, which interpolates geometry attributes at each time step.
Static runtime parameters fixed for the simulation, and will trigger a recompilation if changed.
A dynamic runtime parameters provider, which interpolates dynamic runtime parameters at each time step.
Stepper (via the step_fn argument), the solver method.
TimeStepCalculator (via the step_fn argument)
TransportModel (via the step_fn argument)
PedestalModel (via the step_fn argument)

The following sections cover these inputs.

Note that, while the standard flow detailed above uses the “config”, and then builds a torax.sim.Sim object, these steps are optional. Users who wish to customize TORAX with purpose-built transport models or steppers may call run_simulation() directly, or use the API to build customized timeloops and not run `run_simulation() at all.

TORAX state

The torax.state module describes the input and output states of TORAX. torax.state.ToraxSimState is the complete simulation state, with several attributes inspired by the IMAS schema. Each time step of a TORAX simulation updates that state.

Geometry

Geometry describes the geometry of the torus and fields. See the torax.geometry package for different ways to build a Geometry and its child classes.

Each Geometry object represents the geometry at a single point in time. A sequence of different geometries at different timesteps can be provided at config, and the various Geometry``s are constructed upon initialization. These are packaged together into a ``GeometryProvider which interpolates a new Geometry at each timestep. This is needed in general, since the TORAX timesteps are not necessarily deterministic.

Config and Runtime Parameters

Within TORAX, we refer to a “config” as a Python file or dictionary configuring the building blocks of a TORAX simulation run. A “config” may define which time-step calculator to use, which transport model to load, and where the geometry comes from.

We refer to “runtime parameters” as the inputs to the various objects running within a TORAX simulation. These runtime parameters may be dynamic, meaning potentially time-varying and will not trigger a JAX recompilation if they change. Or they are static, meaning that they will trigger a recompilation if they change from run to run.

The torax.config module contains files for both of these.

torax.config.build_sim builds a torax.sim.Sim object from a config dictionary, like the ones see in the torax/examples/ folder.

torax.config.runtime_params shows the runtime inputs to the TORAX simulation, and torax.interpolated_param shows the logic of how we take user input configs and interpolate them to a specific time-step, allowing for time-varying runtime params.

Many of the modules below also have module-specific runtime parameters that are also fed into the simulation. All these runtime params are packaged together and interpolated to a specific time, referred to as a params “slice”, torax.config.runtime_params_slice.

Stepper

torax.stepper contains PDE time steppers that discretize the PDE in time and solve for the next time step with linear or nonlinear methods.

Inside the Stepper implementations is where JAX is actually used to compute Jacobians or do optimization-based solving. See the implementations for more details.

Most Steppers are built with SourceModels , a TransportModel, and a PedestalModel, described below.

The Stepper class is abstract and can be extended. Users may provide their own implementation and feed it to torax.sim.run_simulation().

Sources

The torax.sources module contains all source models plugged into TORAX. They are packaged together into a SourceModels object, which is a simple container to help access all the sources while stepping through the simulation.

A TORAX Source produces heat, particle, or current deposition profiles used to compute PDE source/sink coefficients used while solving for the next simulation state. TORAX provides several default source model implementations, all of which are configurable via the Python dict config, but users may also extend Source and add their own.

More details on how to create new sources in How to integrate new models.

Transport model

A TORAX TransportModel computes the heat and particle turbulent transport coefficients. TransportModel is an abstract class, and TORAX provides several implementations, including QLKNN.

See the torax.transport_model module for all implementations. Users may extend TransportModel to create their own implementation as well. More details in How to integrate new models.

Pedestal model

A TORAX PedestalModel imposes the plasma temperature and density at a desired internal location. This is intended to correspond to the top of the H-mode pedstal. The operation of the pedestal is controlled by a time-dependent configuration attribute. PedestalModel is an abstract class, and TORAX currently provides two simple implementations.

See the torax.pedestal_model module for all implementations. Users may extend PedestalModel to create their own implementation as well.

Time step calculator

torax.time_step_calculator contains the interface and default implementations of TimeStepCalculator, the base class which computes the duration of the next time step in TORAX and decides when the simulation is over.

Users may use one of the provided implementations or create their own by extending TimeStepCalculator.