From 058d0f0cca05ac212b91a037aabba06ddac49bf2 Mon Sep 17 00:00:00 2001 From: Michael Zingale Date: Fri, 8 Dec 2023 12:46:33 -0500 Subject: [PATCH] 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 --- .../analysis/massive_star_multi.py | 60 +++++++++++++------ 1 file changed, 41 insertions(+), 19 deletions(-) diff --git a/Exec/science/massive_star/analysis/massive_star_multi.py b/Exec/science/massive_star/analysis/massive_star_multi.py index 08f4e9a25f..29f6bf3829 100755 --- a/Exec/science/massive_star/analysis/massive_star_multi.py +++ b/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)