add a bit more flexibility to the massive star plotting script (#2671)

This can now do a vertical orientation and also pass in the width for zooming via the commandline

Exec/science/massive_star/analysis/massive_star_multi.py

@@ -1,22 +1,24 @@
 #!/usr/bin/env python3
 
+import argparse
+import os
+
 import matplotlib
-matplotlib.use('agg')
+import matplotlib.pyplot as plt
 
-import os
-import sys
 import yt
-import matplotlib.pyplot as plt
-import numpy as np
-from functools import reduce
+from yt.frontends.boxlib.api import CastroDataset
+# assume that our data is in CGS
+from yt.units import cm
+
+matplotlib.use('agg')
 
-from mpl_toolkits.axes_grid1 import ImageGrid
 
-# assume that our data is in CGS
-from yt.units import cm, amu
-from yt.frontends.boxlib.api import CastroDataset
 
-def make_plot(plotfile, prefix="plot", size=(19.2, 10.8), cbar_location="right"):
+
+def make_plot(plotfile, prefix="plot", width_frac=0.1,
+              layout=(1, 4),
+              size=(19.2, 10.8), cbar_location="right"):
 
     ds = CastroDataset(plotfile)
 
@@ -36,8 +38,6 @@ def make_plot(plotfile, prefix="plot", size=(19.2, 10.8), cbar_location="right")
     fig = plt.figure()
     fig.set_size_inches(size)
 
-    width_frac = 0.1
-
     fields = ["MachNumber", "magvort", "abar", "enuc"]
 
     sp = yt.SlicePlot(ds, "theta", fields,
@@ -75,17 +75,39 @@ def make_plot(plotfile, prefix="plot", size=(19.2, 10.8), cbar_location="right")
 
     sp.set_axes_unit("cm")
 
-    fig = sp.export_to_mpl_figure((1, len(fields)), cbar_location=cbar_location, cbar_pad="5%")
+    fig = sp.export_to_mpl_figure((layout[0], layout[1]), axes_pad=(1.0, 0.4),
+                                  cbar_location=cbar_location, cbar_pad="2%")
 
     fig.subplots_adjust(left=0.05, right=0.95, bottom=0.025, top=0.975)
-    fig.text(0.02, 0.02, "time = {:8.5f} s".format(float(ds.current_time)), transform=fig.transFigure)
+    fig.text(0.02, 0.02, f"time = {float(ds.current_time):8.2f} s",
+             transform=fig.transFigure)
     fig.set_size_inches(size)
+    fig.tight_layout()
+    extra = ""
+    if layout[0] >= layout[1]:
+        extra = "_vertical"
 
-    fig.savefig(f"{prefix}_{os.path.basename(plotfile)}_slice.png", pad_inches=0)
+    fig.savefig(f"{prefix}_{os.path.basename(plotfile)}_w{width_frac:04.2f}{extra}.pdf", pad_inches=0)
 
 
 if __name__ == "__main__":
 
-    plotfile = sys.argv[1]
-
-    make_plot(plotfile, "all")
+    p = argparse.ArgumentParser()
+    p.add_argument("--vertical", action="store_true",
+                   help="plot 2x2 or 1x4")
+    p.add_argument("--width_fraction", type=float, default=0.1,
+                   help="fraction of domain to show")
+    p.add_argument("plotfile", type=str, nargs=1,
+                   help="plotfile to plot")
+
+    args = p.parse_args()
+    plotfile = args.plotfile[0]
+
+    if args.vertical:
+        size = (7.5, 11.0)
+        layout = (2, 2)
+    else:
+        size = (19.2, 8.5)
+        layout = (1, 4)
+
+    make_plot(plotfile, "all", layout=layout, width_frac=args.width_fraction, size=size)