Fix plotting script following new xarray output.

README.md

@@ -199,7 +199,7 @@ python3 run_simulation_main.py \
 
 ### Post-simulation
 
-Once complete, the time history of a simulation state and derived quantities is written to `state_history.h5`. The output path is written to stdout
+Once complete, the time history of a simulation state and derived quantities is written to `state_history.nc`. The output path is written to stdout
 
 To take advantage of the in-memory (non-persistent) cache, the process does not end upon simulation termination. It is possible to modify the config, toggle the `log_progress` and `plot_progress` flags, and rerun the simulation. Only the following modifications will then trigger a recompilation:
 
@@ -219,21 +219,3 @@ deactivate
 # Simulation tutorials
 
 Under construction
-
-# FAQ
-
-* On MacOS, you may get the error: .. ERROR:: Could not find a local HDF5
-  installation.:
-* Solution: You need to tell the OS where HDF5 is, try
-
-```shell
-brew install hdf5
-```
-
-```shell
-export HDF5_DIR="$(brew --prefix hdf5)"
-```
-
-```shell
-pip install --no-binary=h5py h5py
-```

requirements/requirements.txt

@@ -11,6 +11,5 @@ setuptools;python_version>="3.10"
 chex>=0.1.85
 fancyflags>=1.2
 equinox @ git+https://github.com/patrick-kidger/equinox@1e601672d38d2c4d483535070a3572d8e8508a20
-h5py>=3.10.0
 PyYAML>=6.0.1
 xarray>=2023.12.0

torax/plotting/plotruns.py

@@ -14,7 +14,7 @@
 
 """Basic post-run plotting tool. Plot a single run or comparison of two runs.
 
-Includes a time slider. Reads output h5 files,
+Includes a time slider. Reads output files with xarray data or legacy h5 data.
 
 Plots:
 (1) chi_i, chi_e (transport coefficients)
@@ -27,11 +27,11 @@
 
 import argparse
 import dataclasses
-import h5py
 import matplotlib
 import matplotlib.pyplot as plt
 from matplotlib.widgets import Slider  # pylint: disable=g-importing-member
 import numpy as np
+import xarray as xr
 
 matplotlib.use('TkAgg')
 
@@ -77,7 +77,7 @@ def __post_init__(self):
     '--outfile',
     nargs='*',
     help=(
-        'Relative location of output h5 files (if two are provided, a'
+        'Relative location of output files (if two are provided, a'
         ' comparison is done)'
     ),
 )
@@ -94,42 +94,50 @@ def __post_init__(self):
 if not args.outfile:
   raise ValueError('No output file provided')
 
-with h5py.File(args.outfile[0] + '.h5', 'r') as hf:
-  plotdata1 = PlotData(
-      ti=hf['temp_ion'][:],
-      te=hf['temp_el'][:],
-      ne=hf['ne'][:],
-      j=hf['jtot'][:],
-      johm=hf['johm'][:],
-      j_bootstrap=hf['j_bootstrap'][:],
-      jext=hf['jext'][:],
-      q=hf['q_face'][:],
-      s=hf['s_face'][:],
-      chi_i=hf['chi_face_ion'][:],
-      chi_e=hf['chi_face_el'][:],
-      t=hf['t'][:],
-      rcell_coord=hf['r_cell_norm'][:],
-      rface_coord=hf['r_face_norm'][:],
-  )
+ds1 = xr.open_dataset(args.outfile[0])
+if 'time' in ds1:
+  t = ds1['time'].to_numpy()
+else:
+  t = ds1['t'].to_numpy()
+plotdata1 = PlotData(
+    ti=ds1['temp_ion'].to_numpy(),
+    te=ds1['temp_el'].to_numpy(),
+    ne=ds1['ne'].to_numpy(),
+    j=ds1['jtot'].to_numpy(),
+    johm=ds1['johm'].to_numpy(),
+    j_bootstrap=ds1['j_bootstrap'].to_numpy(),
+    jext=ds1['jext'].to_numpy(),
+    q=ds1['q_face'].to_numpy(),
+    s=ds1['s_face'].to_numpy(),
+    chi_i=ds1['chi_face_ion'].to_numpy(),
+    chi_e=ds1['chi_face_el'].to_numpy(),
+    rcell_coord=ds1['r_cell_norm'].to_numpy(),
+    rface_coord=ds1['r_face_norm'].to_numpy(),
+    t=t,
+)
 
 if comp_plot:
-  with h5py.File(args.outfile[1] + '.h5', 'r') as hf:
-    plotdata2 = PlotData(
-        ti=hf['temp_ion'][:],
-        te=hf['temp_el'][:],
-        ne=hf['ne'][:],
-        j=hf['jtot'][:],
-        johm=hf['johm'][:],
-        j_bootstrap=hf['j_bootstrap'][:],
-        jext=hf['jext'][:],
-        q=hf['q_face'][:],
-        s=hf['s_face'][:],
-        chi_i=hf['chi_face_ion'][:],
-        chi_e=hf['chi_face_el'][:],
-        t=hf['t'][:],
-        rcell_coord=hf['r_cell_norm'][:],
-        rface_coord=hf['r_face_norm'][:],
-    )
+  ds2 = xr.open_dataset(args.outfile[1])
+  if 'time' in ds2:
+    t = ds2['time'].to_numpy()
+  else:
+    t = ds2['t'].to_numpy()
+  plotdata2 = PlotData(
+      ti=ds2['temp_ion'].to_numpy(),
+      te=ds2['temp_el'].to_numpy(),
+      ne=ds2['ne'].to_numpy(),
+      j=ds2['jtot'].to_numpy(),
+      johm=ds2['johm'].to_numpy(),
+      j_bootstrap=ds2['j_bootstrap'].to_numpy(),
+      jext=ds2['jext'].to_numpy(),
+      q=ds2['q_face'].to_numpy(),
+      s=ds2['s_face'].to_numpy(),
+      chi_i=ds2['chi_face_ion'].to_numpy(),
+      chi_e=ds2['chi_face_el'].to_numpy(),
+      rcell_coord=ds2['r_cell_norm'].to_numpy(),
+      rface_coord=ds2['r_face_norm'].to_numpy(),
+      t=t,
+  )
 
 fig = plt.figure(figsize=(15, 10))
 ax1 = fig.add_subplot(231)

torax/simulation_app.py

@@ -75,7 +75,7 @@ class AnsiColors(enum.Enum):
 _ANSI_END = '\033[0m'
 
 _DEFAULT_OUTPUT_DIR_PREFIX = '/tmp/torax_results_'
-_STATE_HISTORY_FILENAME = 'state_history.h5'
+_STATE_HISTORY_FILENAME = 'state_history.nc'
 
 
 def log_to_stdout(output: str, color: AnsiColors | None = None) -> None:

torax/tests/test_lib/sim_test_case.py

@@ -22,7 +22,6 @@
 from absl.testing import absltest
 from absl.testing import parameterized
 import chex
-import h5py
 import jax.numpy as jnp
 import numpy as np
 import torax
@@ -36,17 +35,7 @@
 from torax.stepper import stepper as stepper_lib
 from torax.time_step_calculator import array_time_step_calculator
 from torax.transport_model import transport_model as transport_model_lib
-
-
-_TORAX_TO_PINT = {
-    'temp_ion': 'Ti',
-    'temp_el': 'Te',
-    'psi': 'psi',
-    's_face': 's',
-    'q_face': 'q',
-    'ne': 'ne',
-}
-
+import xarray as xr
 
 _PYTHON_MODULE_PREFIX = '.tests.test_data.'
 _PYTHON_CONFIG_PACKAGE = 'torax'
@@ -116,19 +105,15 @@ def _get_refs(
     """Gets reference values for the requested state profiles."""
     expected_results_path = self._expected_results_path(ref_name)
     self.assertTrue(os.path.exists(expected_results_path))
-
-    with open(expected_results_path, mode='rb') as f:
-      with h5py.File(f, 'r') as hf:
-        self.assertNotEmpty(profiles)
-        if 'Ti' in hf.keys():  # Determine if h5 file is PINT output
-          ref_profiles = {
-              profile: hf[_TORAX_TO_PINT[profile]][:] for profile in profiles
-          }
-        else:
-          ref_profiles = {profile: hf[profile][:] for profile in profiles}
-        ref_time = jnp.array(hf['t'])
-        self.assertEqual(ref_time.shape[0], ref_profiles[profiles[0]].shape[0])
-        return ref_profiles, ref_time
+    ds = xr.open_dataset(expected_results_path)
+    self.assertNotEmpty(profiles)
+    ref_profiles = {profile: ds[profile].to_numpy() for profile in profiles}
+    if 'time' in ds:
+      ref_time = ds['time'].to_numpy()
+    else:
+      ref_time = ds['t'].to_numpy()
+    self.assertEqual(ref_time.shape[0], ref_profiles[profiles[0]].shape[0])
+    return ref_profiles, ref_time
 
   def _check_profiles_vs_expected(
       self,