Running simulations programmatically

This short guide describes how to integrate Torax into your codebase and allows you to run it multiple times efficiently.

First, we need a torax.ToraxConfig object representing the simulation config. In this example, we will use the torax/examples/iterhybrid_rampup.py config:

import torax

torax_config = torax.build_torax_config_from_file('examples/iterhybrid_rampup.py')

If you already have a config_dict dictionary in Python, you could instead use torax_config = torax.ToraxConfig.from_dict(config_dict).

We can then run the simulation:

# returns the output XArray DataTree and a torax.StateHistory object.
data_tree, state_history = torax.run_simulation(torax_config)

# Check that the simulation completed successfully.
if state_history.sim_error != torax.SimError.NO_ERROR:
  raise ValueError(
      f'TORAX failed to run the simulation with error: {state_history.sim_error}.'
  )

# Example below shows how to access the fusion gain at time=2 seconds.
Q_fusion_t2 = data_tree.scalars.Q_fusion.sel(time=2, method='nearest')

Plotting from an in-memory simulation

If you have already run a simulation and have a data_tree in memory, you can plot it directly without saving to a file first using torax.plot_run_from_data_tree:

plot_config = torax.import_module('plotting/configs/default_plot_config.py')['PLOT_CONFIG']

# Plot directly from the in-memory data_tree returned by run_simulation.
fig = torax.plot_run_from_data_tree(plot_config, data_tree)

To compare two in-memory runs:

fig = torax.plot_run_from_data_tree(plot_config, data_tree, data_tree2)

If you have saved the output to a .nc file and want to plot from disk, use torax.plot_run instead:

fig = torax.plot_run(plot_config, PATH_TO_LOCAL_NC_FILE)