Plot with datatree w/o creating .nc file

PiperOrigin-RevId: 884592463

Author: skdeepmind

docs/plotting.rst

@@ -175,13 +175,16 @@ plot layout and content. Dataclass fields and defaults are as follows:
 - ``cols`` (int): Number of columns in the figure.
 - ``axes`` (tuple of ``PlotProperties``):  Configuration for each subplot.
   See below.
-- ``figure_size_factor`` (float=5.0): Scaling factor for the figure size.
-- ``tick_fontsize`` (int=10): Font size for axis ticks.
-- ``axes_fontsize`` (int=10): Font size for axis labels.
-- ``title_fontsize`` (int=16): Font size for the figure title.
-- ``default_legend_fontsize`` (int=10): Default font size for legends.
-- ``colors`` (tuple[str, ...] = ('r', 'b', 'g', 'm', 'y', 'c')): Colors to use
-  for plot lines. Cycles through the tuple for multiple lines.
+- ``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

docs/running_programmatically.rst

@@ -33,3 +33,30 @@ We can then run the simulation:
 
   # 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``:
+
+.. code-block:: python
+
+  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:
+
+.. code-block:: python
+
+  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:
+
+.. code-block:: python
+
+  fig = torax.plot_run(plot_config, PATH_TO_LOCAL_NC_FILE)

torax/__init__.py

@@ -32,6 +32,7 @@
 from torax._src.output_tools.post_processing import PostProcessedOutputs
 from torax._src.pedestal_model.pedestal_model_output import PedestalModelOutput
 from torax._src.plotting.plotruns_lib import plot_run
+from torax._src.plotting.plotruns_lib import plot_run_from_data_tree
 from torax._src.sources.source_profiles import SourceProfiles
 from torax._src.state import CoreProfiles
 from torax._src.state import CoreTransport
@@ -58,6 +59,7 @@
     'build_torax_config_from_file',
     'import_module',
     'plot_run',
+    'plot_run_from_data_tree',
     'run_simulation',
     'CoreProfiles',
     'CoreTransport',

torax/_src/plotting/plotruns_lib.py

@@ -15,7 +15,8 @@
 """Utilities for plotting outputs of Torax runs.
 
 Public API:
-  plot_run: Main entry point. Loads data and returns a plotly Figure.
+  plot_run: Main entry point. Loads data from file and returns a plotly Figure.
+  plot_run_from_data_tree: Plots from an in-memory xr.DataTree.
   PlotData: Data container exposing all output variables as attributes.
   FigureProperties: Configuration for the overall figure layout.
   PlotProperties: Configuration for an individual subplot.
@@ -315,62 +316,67 @@ def available_variables(self) -> set[str]:
     return attrs
 
 
-def load_data(filename: str) -> PlotData:
-  """Loads an xr.Dataset from a file, handling coordinate name changes."""
-
-  data_tree = output.load_state_file(filename)
+def _transform_data(ds: xr.Dataset) -> xr.Dataset:
+  """Transforms dataset variables to plotting units."""
+  # TODO(b/414755419)
+  ds = ds.copy()
+
+  transformations = {
+      output.J_TOROIDAL_TOTAL: 1e6,  # A/m^2 to MA/m^2
+      output.J_TOROIDAL_OHMIC: 1e6,  # A/m^2 to MA/m^2
+      output.J_TOROIDAL_BOOTSTRAP: 1e6,  # A/m^2 to MA/m^2
+      output.J_TOROIDAL_EXTERNAL: 1e6,  # A/m^2 to MA/m^2
+      'j_generic_current': 1e6,  # A/m^2 to MA/m^2
+      output.I_BOOTSTRAP: 1e6,  # A to MA
+      output.IP_PROFILE: 1e6,  # A to MA
+      output.IP: 1e6,  # A to MA
+      'j_ecrh': 1e6,  # A/m^2 to MA/m^2
+      'p_icrh_i': 1e6,  # W/m^3 to MW/m^3
+      'p_icrh_e': 1e6,  # W/m^3 to MW/m^3
+      'p_generic_heat_i': 1e6,  # W/m^3 to MW/m^3
+      'p_generic_heat_e': 1e6,  # W/m^3 to MW/m^3
+      'p_ecrh_e': 1e6,  # W/m^3 to MW/m^3
+      'p_alpha_i': 1e6,  # W/m^3 to MW/m^3
+      'p_alpha_e': 1e6,  # W/m^3 to MW/m^3
+      'p_ohmic_e': 1e6,  # W/m^3 to MW/m^3
+      'p_bremsstrahlung_e': 1e6,  # W/m^3 to MW/m^3
+      'p_cyclotron_radiation_e': 1e6,  # W/m^3 to MW/m^3
+      'p_impurity_radiation_e': 1e6,  # W/m^3 to MW/m^3
+      'ei_exchange': 1e6,  # W/m^3 to MW/m^3
+      'P_ohmic_e': 1e6,  # W to MW
+      'P_aux_total': 1e6,  # W to MW
+      'P_alpha_total': 1e6,  # W to MW
+      'P_bremsstrahlung_e': 1e6,  # W to MW
+      'P_cyclotron_e': 1e6,  # W to MW
+      'P_ecrh': 1e6,  # W to MW
+      'P_radiation_e': 1e6,  # W to MW
+      'I_ecrh': 1e6,  # A to MA
+      'I_aux_generic': 1e6,  # A to MA
+      'W_thermal_total': 1e6,  # J to MJ
+      output.N_E: 1e20,  # m^-3 to 10^{20} m^-3
+      output.N_I: 1e20,  # m^-3 to 10^{20} m^-3
+      output.N_IMPURITY: 1e20,  # m^-3 to 10^{20} m^-3
+  }
 
-  def _transform_data(ds: xr.Dataset):
-    """Transforms data in-place to the desired units."""
-    # TODO(b/414755419)
-    ds = ds.copy()
-
-    transformations = {
-        output.J_TOROIDAL_TOTAL: 1e6,  # A/m^2 to MA/m^2
-        output.J_TOROIDAL_OHMIC: 1e6,  # A/m^2 to MA/m^2
-        output.J_TOROIDAL_BOOTSTRAP: 1e6,  # A/m^2 to MA/m^2
-        output.J_TOROIDAL_EXTERNAL: 1e6,  # A/m^2 to MA/m^2
-        'j_generic_current': 1e6,  # A/m^2 to MA/m^2
-        output.I_BOOTSTRAP: 1e6,  # A to MA
-        output.IP_PROFILE: 1e6,  # A to MA
-        output.IP: 1e6,  # A to MA
-        'j_ecrh': 1e6,  # A/m^2 to MA/m^2
-        'p_icrh_i': 1e6,  # W/m^3 to MW/m^3
-        'p_icrh_e': 1e6,  # W/m^3 to MW/m^3
-        'p_generic_heat_i': 1e6,  # W/m^3 to MW/m^3
-        'p_generic_heat_e': 1e6,  # W/m^3 to MW/m^3
-        'p_ecrh_e': 1e6,  # W/m^3 to MW/m^3
-        'p_alpha_i': 1e6,  # W/m^3 to MW/m^3
-        'p_alpha_e': 1e6,  # W/m^3 to MW/m^3
-        'p_ohmic_e': 1e6,  # W/m^3 to MW/m^3
-        'p_bremsstrahlung_e': 1e6,  # W/m^3 to MW/m^3
-        'p_cyclotron_radiation_e': 1e6,  # W/m^3 to MW/m^3
-        'p_impurity_radiation_e': 1e6,  # W/m^3 to MW/m^3
-        'ei_exchange': 1e6,  # W/m^3 to MW/m^3
-        'P_ohmic_e': 1e6,  # W to MW
-        'P_aux_total': 1e6,  # W to MW
-        'P_alpha_total': 1e6,  # W to MW
-        'P_bremsstrahlung_e': 1e6,  # W to MW
-        'P_cyclotron_e': 1e6,  # W to MW
-        'P_ecrh': 1e6,  # W to MW
-        'P_radiation_e': 1e6,  # W to MW
-        'I_ecrh': 1e6,  # A to MA
-        'I_aux_generic': 1e6,  # A to MA
-        'W_thermal_total': 1e6,  # J to MJ
-        output.N_E: 1e20,  # m^-3 to 10^{20} m^-3
-        output.N_I: 1e20,  # m^-3 to 10^{20} m^-3
-        output.N_IMPURITY: 1e20,  # m^-3 to 10^{20} m^-3
-    }
+  for var_name, scale in transformations.items():
+    if var_name in ds:
+      ds[var_name] /= scale
 
-    for var_name, scale in transformations.items():
-      if var_name in ds:
-        ds[var_name] /= scale
+  return ds
 
-    return ds
 
+def _data_tree_to_plot_data(data_tree: xr.DataTree) -> PlotData:
+  """Converts an xr.DataTree to a PlotData object with unit transformations."""
   return PlotData(xr.map_over_datasets(_transform_data, data_tree))
 
 
+def load_data(filename: str) -> PlotData:
+  """Loads an xr.Dataset from a file, handling coordinate name changes."""
+
+  data_tree = output.load_state_file(filename)
+  return _data_tree_to_plot_data(data_tree)
+
+
 def _get_file_path(outfile: str) -> str:
   """Gets the absolute path to the file."""
   possible_paths = [outfile]
@@ -399,23 +405,38 @@ def plot_run(
     outfile2: str | None = None,
     interactive: bool = True,
 ) -> go.Figure:
-  """Plots a single run or comparison of two runs."""
+  """Plots a single run or comparison of two runs from output files."""
+  fig_title = plot_config.figure_title or _get_title_from_paths(
+      outfile, outfile2
+  )
+
   outfile = _get_file_path(outfile)
   outfile2 = _get_file_path(outfile2) if outfile2 else None
 
-  plotdata1 = load_data(outfile)
-  plotdata2 = load_data(outfile2) if outfile2 else None
+  data_tree = output.load_state_file(outfile)
+  data_tree2 = output.load_state_file(outfile2) if outfile2 else None
+  return plot_run_from_data_tree(
+      plot_config, data_tree, data_tree2, interactive, fig_title
+  )
+
 
-  # Prepare list of datasets to check, associating them with their filenames
-  # for clearer errors
-  datasets_to_check = [(plotdata1, outfile)]
+def plot_run_from_data_tree(
+    plot_config: FigureProperties,
+    data_tree: xr.DataTree,
+    data_tree2: xr.DataTree | None = None,
+    interactive: bool = True,
+    fig_title: str = 'Torax Simulation Results',
+) -> go.Figure:
+  """Plots a single run or comparison of two runs from in-memory DataTrees."""
+  plotdata1 = _data_tree_to_plot_data(data_tree)
+  plotdata2 = _data_tree_to_plot_data(data_tree2) if data_tree2 else None
+
+  datasets_to_check = [(plotdata1, 'data_tree')]
   if plotdata2 is not None:
-    datasets_to_check.append((plotdata2, outfile2))
+    datasets_to_check.append((plotdata2, 'data_tree2'))
 
   for plotdata, filename in datasets_to_check:
-    # Get the set of valid keys for this specific dataset
     available_vars = plotdata.available_variables()
-
     for cfg in plot_config.axes:
       for attr in cfg.attrs:
         if attr not in available_vars:
@@ -424,8 +445,7 @@ def plot_run(
               f'output file: {filename}'
           )
 
-  title = plot_config.figure_title or _get_title_from_paths(outfile, outfile2)
-  fig = create_plotly_figure(plot_config, plotdata1, plotdata2, title)
+  fig = create_plotly_figure(plot_config, plotdata1, plotdata2, fig_title)
   if interactive:
     fig.show()