Add fast ions to outputs.

PiperOrigin-RevId: 879024158

Author: hamelphi

torax/_src/output_tools/output.py

@@ -654,6 +654,22 @@ def _save_core_profiles(
         coords={MAIN_ION: main_ions, TIME: self.times},
         name="main_ion_fractions",
     )
+    # Handle fast ions
+    first_fast_ions = self.core_profiles[0].fast_ions
+    for i, first_fi in enumerate(first_fast_ions):
+      source_key = f"{first_fi.source}_{first_fi.species}"
+      n_data = np.stack([
+          cp.fast_ions[i].n.cell_plus_boundaries() for cp in self.core_profiles
+      ])
+      T_data = np.stack([
+          cp.fast_ions[i].T.cell_plus_boundaries() for cp in self.core_profiles
+      ])
+      xr_dict[f"n_fast_ion_{source_key}"] = self._pack_into_data_array(
+          f"n_fast_ion_{source_key}", n_data
+      )
+      xr_dict[f"T_fast_ion_{source_key}"] = self._pack_into_data_array(
+          f"T_fast_ion_{source_key}", T_data
+      )
 
     return xr_dict