Plotting simulations

TORAX provides plotting functionalities to visualize simulation outputs and compare different runs. The plots include a time slider for interactive manipulation.

Using the plot_torax Script

To visualize simulation results, use the plot_torax script. This script offers flexibility in plotting single runs, comparing two runs, and customizing plot configurations.

Plotting a Single Run

To plot the output of a single simulation, run the following command from the TORAX root directory:

plot_torax --outfile <full_path_to_simulation_output> \
 --plot_config <path_to_plot_config>

Replace <full_path_to_simulation_output> with the full path to your simulation’s output file. Optionally, specify a custom plot configuration module using --plot_config, with the path for the Python file containing the PLOT_CONFIG variable. See Customizable Plot Configurations for more details.

Comparing Two Runs

To compare the output of two simulations, provide both output file paths:

plot_torax --outfile <full_path_to_simulation_output1> \
 <full_path_to_simulation_output2>

This overlays the plots of the two runs, allowing for direct comparison.

Customizable Plot Configurations

The --plot_config flag allows you to define which quantities are plotted and how they are arranged. This flag accepts the path to a Python file containing a PLOT_CONFIG variable which is an instance of torax.plotting.plotruns_lib.FigureProperties. This variable specifies the layout and properties of the plot. If no --plot_config is specified, the default configuration |default_plot_config.py| is used.

v1.4.0 is the version of the Torax codebase.

Prepackaged Plot Configurations

Several pre-packaged plot configurations are available in the plotting/configs directory. These include:

default_plot_config.py - A default configuration with various profiles and scalars including kinetic profiles, heating and current source profiles, transport coefficients, and post-processed output quantities.

simple_plot_config.py - A simplified configuration containing a subset of the default configuration. Plots kinetic profiles, currents, transport coefficients, safety factor and magnetic shear.

sources_plot_config.py - A configuration focusing on profile and scalar plots of heat, particle and current source profiles and their integrated quantities.

global_params_plot_config.py - A configuration focusing on scalar parameters and their evolution over time, to provide a quick overview of the simulation. Includes total currents, fusion gain, averaged kinetic profiles, q95.

Example usage:

plot_torax --outfile <full_path_to_simulation_output> \
 --plot_config plotting/configs/sources_plot_config.py

These plot configurations can also be passed as input into the run_torax script using the --plot_config flag. See run_torax flags.

The PLOT_CONFIG from plotting/configs/simple_plot_config.py is shown below as an example.

from torax._src.plotting import plotruns_lib

PLOT_CONFIG = plotruns_lib.FigureProperties(
    rows=2,
    cols=3,
    axes=(
        plotruns_lib.PlotProperties(
            attrs=('T_i', 'T_e'),
            labels=(r'$T_i$', r'$T_e$'),
            ylabel='Temperature [keV]',
        ),
        plotruns_lib.PlotProperties(
            attrs=('n_e',),
            labels=(r'$n_e$',),
            ylabel=r'Electron density $[10^{20}~m^{-3}]$',
        ),
        plotruns_lib.PlotProperties(
            attrs=('chi_turb_i', 'chi_turb_e'),
            labels=(r'$\chi_i$', r'$\chi_e$'),
            ylabel=r'Heat conductivity $[m^2/s]$',
            upper_percentile=98.0,
            include_first_timepoint=False,
            ylim_min_zero=False,
        ),
        plotruns_lib.PlotProperties(
            attrs=(
                'j_total',
                'j_ohmic',
                'j_bootstrap',
                'j_generic_current',
                'j_ecrh',
            ),
            labels=(
                r'$j_{tot}$',
                r'$j_{ohm}$',
                r'$j_{bs}$',
                r'$j_{generic}$',
                r'$j_{ecrh}$',
            ),
            ylabel=r'Toroidal current $[MA~m^{-2}]$',
            legend_fontsize=8,
        ),
        plotruns_lib.PlotProperties(
            attrs=('q',),
            labels=(r'$q$',),
            ylabel='Safety factor',
        ),
        plotruns_lib.PlotProperties(
            attrs=('magnetic_shear',),
            labels=(r'$\hat{s}$',),
            ylabel='Magnetic shear',
        ),
    ),
)

Customizing Plots

Customizing plots is not currently exposed as part of the public API. If you would like to customize plots below is detailed how you can do so for Torax version 1.0.0 but this is not guaranteed to be supported in development.

In the near future we aim to provide a stable and supported API for customizing plots.

For now there are three objects that are useful for customizing plots:

from torax._src.plotting.plotruns_lib import FigureProperties
from torax._src.plotting.plotruns_lib import PlotProperties
from torax._src.plotting.plotruns_lib import PlotType

The FigureProperties dataclass offers several options for customizing the plot layout and content. Dataclass fields and defaults are as follows:

  • rows (int): Number of rows in the figure.

  • cols (int): Number of columns in the figure.

  • axes (tuple of PlotProperties): Configuration for each subplot. See below.

  • font_family (str = ‘Arial, sans-serif’): Font family for all text in the figure.

  • title_size (int = 16): Font size for the main figure title.

  • subplot_title_size (int = 12): Font size for subplot titles.

  • tick_size (int = 8): Font size for axis ticks.

  • nticks_time (int = 6): Number of x-axis ticks for the time series plots.

  • tickvals_rho (tuple[float, …] = (0, 0.2, 0.4, 0.6, 0.8, 1.0)): Values of x-axis ticks for the spatial plots.

  • height (int | None = None): Height of the figure in pixels. If None, autosize is used.

  • legend_spacing (int = 10): Spacing between legend entries.

  • margin (dict[str, int] = {‘l’: 40, ‘r’: 40, ‘t’: 80, ‘b’: 40}): Margin around the figure. A dict with keys l, r, t and b for left, right, top and bottom margins respectively.

The PlotProperties dataclass configures individual subplots. For example, the PlotProperties object for plotting ion and electron temperatures looks like this:

plotruns_lib.PlotProperties(
    attrs=('T_i', 'T_e'),
    labels=(r'$T_i$', r'$T_e$'),
    ylabel='Temperature [keV]',
),

The fields in PlotProperties are as follows:

  • attrs: Tuple of attribute names from the PlotData dataclass used to retrieve the data for plotting.

  • labels: Tuple of labels for the plotted lines, one label per attribute in attrs.

  • ylabel: Label for the y-axis.

  • legend_fontsize (int | None): Legend font size. If None, defaults to default_legend_fontsize in FigureProperties.

  • upper_percentile (float=100.0): Filters out outlier data above a given percentile for plotting purposes.

  • lower_percentile (float=0.0): Filters out outlier data below a given percentile for plotting purposes.

  • include_first_timepoint (bool=True): Whether to include the first time point in calculating plot range.

  • ylim_min_zero (bool=True): Whether the plot limits should start from zero.

  • plot_type (PlotType=PlotType.SPATIAL): Defines whether the plot is a spatial profile, or time series plot.

  • suppress_zero_values (bool=False): If True, all-zero-data is not plotted.

suppress_zero_values is useful when defining plots where not all the attrs may be relevant for all runs. For example, if a run does not include a bootstrap current, the j_bootstrap attribute will be all zero. Setting suppress_zero_values=True will automatically exclude this line from the plot.

upper_percentile and lower_percentile are useful for excluding outliers from the plot range calculation, for example transient spikes in the data.

plot_type can be set to either PlotType.SPATIAL (default) or PlotType.TIME_SERIES. Spatial plots are 1D profiles which are updated at each time slice, following time slider manipulation. Time series plots are 0D quantities plotted against the full simulation time, and are not affected by the time slider.

By creating a custom Python module with a PLOT_CONFIG variable set to a FigureProperties instance you can thus completely customize which variables are plotted by Torax by defining a new FigureProperties instance. For examples, see torax/plotting/configs/*.py.

Interactive Time Slider

When plotting with plot_torax, an interactive time slider appears below the plots. This slider allows you to scroll through the simulation output at different timesteps. The spatial profile plots defined above are dynamically updated when the slider is manipulated. The time series plots, defined with plot_type=PlotType.TIME_SERIES stay constant, plotting variables against the full Torax simulation timescale.