Fix field names in diag_parser.py and add example script (#2667)

Make sure np.genfromtxt() doesn't mess with the field names, add an example script that plots mass fractions over time from species_diag.out, and add a note to the docs.


Co-authored-by: Alexander Smith <[email protected]>

Docs/source/io.rst

@@ -562,7 +562,15 @@ By default, 4 output files are created:
 
        The species masses are given in units of solar masses.
 
-Some problems have custom versions of the diagnostics with additional information.
+``Castro/Util/scripts/diag_parser.py`` contains Python code for parsing
+these output files into Numpy arrays.  Usage instructions are included
+in the file, along with an example script at
+``Castro/Util/scripts/plot_species.py``.  This reads a
+``species_diag.out`` file provided on the command line and makes a plot
+of the total mass fractions over time.
+
+Some problems have custom versions of the diagnostics with additional
+information.  These are not currently supported by the Python parser.
 
 
 .. _sec:parallel_io:

Util/scripts/diag_parser.py

@@ -10,7 +10,7 @@
 
 * copy this file into the same directory as your script
 
-Then you can do `from diag_parser import deduplicate, read_diag_file`.
+Then you can do `from diag_parser import read_diag_file`.
 """
 
 from pathlib import Path
@@ -52,10 +52,15 @@
 }
 
 
-def read_diag_file(file_path):
+def read_diag_file(file_path, dedupe=True):
     """Reads a Castro diagnostic file into a numpy structured array.
 
-    Currently only supports the default files that Castro generates.
+    The output can be used directly (the functions in numpy.lib.recfunctions
+    are helpful for this), or converted into a Pandas dataframe (recommended if
+    you plan to do any non-trivial processing).
+
+    Currently only supports the default files that Castro generates, not any of
+    the problem-specific ones.
     """
     if not isinstance(file_path, Path):
         file_path = Path(file_path)
@@ -84,8 +89,16 @@ def read_diag_file(file_path):
                 dtypes[7] = int  # minimum gpu memory free
         # already read the first header line, so we don't need to skip any rows
         data = np.genfromtxt(
-            f, delimiter=widths, comments="#", dtype=dtypes, names=True
+            f,
+            delimiter=widths,
+            comments="#",
+            dtype=dtypes,
+            names=True,
+            deletechars="",
+            replace_space=None,
         )
+    if dedupe:
+        data = deduplicate(data)
     return data
 
 

Util/scripts/plot_species.py

@@ -0,0 +1,36 @@
+#!/usr/bin/env python3
+import sys
+
+import matplotlib as mpl
+import matplotlib.pyplot as plt
+import numpy as np
+from cycler import cycler
+from numpy.lib import recfunctions as rfn
+
+from diag_parser import deduplicate, read_diag_file
+
+data = deduplicate(read_diag_file(sys.argv[1]))
+
+mass_columns = [name for name in rfn.get_names(data.dtype) if name.startswith("Mass ")]
+
+# compute the total mass in the domain
+total_mass = rfn.apply_along_fields(np.sum, data[mass_columns])
+
+# cycle through colors then line styles
+plt.rcParams["axes.prop_cycle"] = (
+    cycler(linestyle=["-", "--"]) * mpl.rcParamsDefault["axes.prop_cycle"]
+)
+
+plt.figure(figsize=(19.20, 10.80))
+for col in mass_columns:
+    plt.plot(data["TIME"], data[col] / total_mass, label=col.removeprefix("Mass "))
+plt.xlabel("Time (s)")
+plt.ylabel("Mass fraction")
+# plt.xscale("log")
+plt.yscale("log")
+
+# put the legend to the right of the plot, not overlapping
+plt.legend(loc="center left", bbox_to_anchor=(1, 0.5))
+
+plt.tight_layout()
+plt.savefig("massfrac_plot.png")