Improvements to the flame_wave flame front analysis scripts

Exec/science/flame_wave/analysis/README.md

@@ -24,20 +24,19 @@ user to specify framerate.
 
 ### front_tracker.py
 Script for measuring the location of a flame or shock front over a sequence of snapshots. Allows the
-user to specify the metrics (will only track 1 / 1000th the local enuc maximum by default), whether
-to use a global or local maximum, domain bounds, and a few other things. For a usage description and
-a full list of valid parameters, type `./front_tracker.py -h`. Should work for any dataset, but has
-only been tested on flame wave ones. Restrictions: Currently only tracks along one dimension (the
-user can tell it how to eliminate the others - either through slicing or averaging), and only tracks
-percentages of the maximum of some field. Outputs to space-delimited data file called
-front_tracking.dat by default.
+user to specify the metrics (will only track 1 / 1000th the enuc maximum by default), domain bounds,
+and a few other things. For a usage description and a full list of valid parameters, type
+`./front_tracker.py -h`. Should work for any dataset, but has only been tested on flame wave ones.
+Restrictions: Currently only tracks along one dimension (the user can tell it how to eliminate the
+others - either through slicing or averaging), and only tracks percentages of the maximum of some
+field. Outputs to space-delimited data file called front_tracking.dat by default.
 
 ### flame_speed.py
-Script for reading in the front tracking dataset, plotting it, and fitting a line to some portion of
-it. Usage: `./flame_speed.py data_file starting_index`, where starting index is the index of the
-first datapoint to consider when fitting the line. The script prints out the slope of the line, the
-r-squared value, and the fit error. The plot will appear squashed when it pops up - the window needs
-to be enlarged. Uses `scipy` and `pandas`.
+Script for reading in the front tracking dataset generated with `front_tracker.py`, plotting it, and
+fitting a line to some portion of it.
+Usage: `./flame_speed.py [--tmin TMIN] [--tmax TMAX] data_file column [column...]`, where `TMIN` and
+`TMAX` specify the times to consider when fitting the line. The script prints out the slope of the
+line, the r-squared value, and the fit error. Uses `scipy` and `pandas`.
 
 ### multirays.py
 Plot 1D vertical slices of axisymmetric datasets. It generates 3 ortho rays - one at each end of

Exec/science/flame_wave/analysis/flame_speed.py

@@ -1,79 +1,156 @@
 #!/usr/bin/env python3
 
-import sys
-import pandas as pd
-from scipy.stats import linregress
+import argparse
+import itertools
+import re
 
-import matplotlib as mpl
 import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+from matplotlib.lines import Line2D
+from matplotlib.patches import Patch
 from matplotlib.ticker import MaxNLocator
+from scipy.stats import linregress
 
-plt.rcParams.update \
-(
+plt.rcParams.update(
     {
-        "font.family": "serif"
+        "font.family": "serif",
     }
 )
 
-def measure_and_plot(dat, stab_ind):
+
+def measure_and_plot(dat, args):
     """
     Plots front_location vs. time on the current pyplot figure, as well as any
-    reasonable linear fits. *stab_ind* should give an index for which the flame
-    front has stabilized, and the slope obtained from linear regression will yield
-    an accurate measurement of the front speed.
+    reasonable linear fits.
     """
+    # pylint: disable=too-many-locals, too-many-statements, too-many-branches
+    # pylint: disable=use-dict-literal
 
-    slopes = dict()
+    slopes = {}
 
-    radarr = dat["enuc[0.1%]"]
     tarr = dat["time"]
-
-    plt.plot(tarr * 1000, radarr, linewidth=2)
-
+    # sort by time
     indx = tarr.argsort()
-    tarr = tarr[indx][stab_ind:]
-    radarr = radarr[indx][stab_ind:]
-
-    m, b, r, _, sd = linregress(tarr, radarr)
-    print("{:>20}\t{:.2e}\t{:.3f}\t{:.2e}".format("enuc[0.1%]", m, r, sd))
-
-    if r > 0.8:
-
-        plt.plot(dat["time"] * 1000, m * dat["time"] + b, "k--", linewidth=2)
+    tarr = tarr[indx]
+
+    common_props = dict(linewidth=2, alpha=0.8)
+
+    # give the same color to <column> and <column>_gmax, regardless of where
+    # they occur in columns
+    color_iter = itertools.cycle(plt.rcParams["axes.prop_cycle"].by_key()["color"])
+    color_lookup = {}
+    any_local_max = False
+    any_global_max = False
+    for column in args.columns:
+        m = re.match(r"(?P<field>.+?)\[(?P<percent>.+?)%\](?P<gmax>_gmax)?", column)
+        assert m is not None
+        field = m["field"]
+        percent = float(m["percent"])
+        global_max = m["gmax"] is not None
+        any_local_max |= not global_max
+        any_global_max |= global_max
+
+        key = (field, percent)
+        if key not in color_lookup:
+            color_lookup[key] = next(color_iter)
+        kwargs = {"color": color_lookup[key]}
+        if not global_max and len(args.columns) > 1:
+            kwargs["linestyle"] = "dashed"
+
+        radarr = dat[column][indx]
+
+        plt.plot(tarr * 1000, radarr, **common_props, **kwargs)
+
+        reg_mask = ~np.isnan(radarr)
+        if args.tmin is not None:
+            reg_mask &= tarr >= args.tmin / 1000
+        if args.tmax is not None:
+            reg_mask &= tarr <= args.tmax / 1000
+
+        if np.count_nonzero(reg_mask) >= 2:
+            m, b, r, _, sd = linregress(tarr[reg_mask], radarr[reg_mask])
+            print(f"{column:<30}\t{m=:.2e}\t{r**2=:.3f}\t{sd=:.2e}")
+
+            if r**2 > 0.8:
+                if len(args.columns) > 1:
+                    color = kwargs["color"]
+                else:
+                    color = "black"
+                plt.plot(
+                    tarr[reg_mask] * 1000,
+                    m * tarr[reg_mask] + b,
+                    **common_props,
+                    color=color,
+                    linestyle="dotted",
+                )
+                slopes[column] = m
 
     ax = plt.gca()
     ax.ticklabel_format(axis="y", style="sci", scilimits=(0, 0))
-    ax.yaxis.major.formatter._useMathText = True
+    ax.yaxis.major.formatter.set_useMathText(True)
     ax.xaxis.set_major_locator(MaxNLocator(integer=True))
     # ax.yaxis.set_major_locator(MaxNLocator(integer=True))
 
+    # build a custom legend
+    line2d_kwargs = []
+    if len(args.columns) > 1:
+        if any_local_max:
+            line2d_kwargs.append(dict(ls="dashed", label="local max"))
+        if any_global_max:
+            line2d_kwargs.append(dict(ls="solid", label="global max"))
+        if slopes:
+            line2d_kwargs.append(dict(ls="dotted", label="linear fit"))
+    legend_elements = [
+        Line2D([0], [0], **common_props, color="black", **kwargs)
+        for kwargs in line2d_kwargs
+    ]
+    legend_elements.extend(
+        Patch(facecolor=color, label=f"{field} {percent:g}%")
+        for (field, percent), color in color_lookup.items()
+    )
+    plt.legend(handles=legend_elements)
+
     plt.xlabel("t [ms]")
     plt.ylabel("r [cm]")
     plt.title("Front Position vs. Time", fontsize=24)
 
     for item in (ax.xaxis.label, ax.yaxis.label):
-
         item.set_fontsize(20)
         item.set_color("dimgrey")
 
     for item in (ax.get_xticklabels() + ax.get_yticklabels()) + [ax.yaxis.offsetText]:
-
         item.set_fontsize(16)
         item.set_color("dimgrey")
 
-    return m
+    return slopes
 
-if __name__ == "__main__":
 
-    try:
+def main():
+    parser = argparse.ArgumentParser(
+        description="Script for plotting and fitting the flame front tracking data generated by front_tracker.py.",
+        epilog="""TMIN and TMAX should give a range of times for which the flame front
+        has stabilized, and the slope obtained from linear regression will
+        yield an accurate measurement of the front speed.""",
+    )
+    parser.add_argument("data_file", type=argparse.FileType("r"))
+    parser.add_argument(
+        "--tmin", type=float, help="minimum time for the linear regression"
+    )
+    parser.add_argument(
+        "--tmax", type=float, help="maximum time for the linear regression"
+    )
+    parser.add_argument("columns", nargs="+", metavar="column")
+    args = parser.parse_args()
 
-        dat = pd.read_csv(sys.argv[1], sep="\s+")
-        stab_ind = int(sys.argv[2])
+    dat = pd.read_csv(args.data_file, sep=r"\s+")
 
-        plt.style.use("ggplot")
-        measure_and_plot(dat, stab_ind)
-        plt.show()
+    plt.style.use("ggplot")
+    plt.figure(figsize=(8, 6))
+    measure_and_plot(dat, args)
+    plt.tight_layout()
+    plt.show()
 
-    except (ValueError, IndexError):
 
-        print("Usage: ./fit_speed.py data_file starting_index")
+if __name__ == "__main__":
+    main()