add a slice plot script

Exec/science/xrb_spherical/analysis/slice.py

@@ -0,0 +1,75 @@
+#!/usr/bin/env python3
+
+import sys
+import os
+import yt
+import numpy as np
+import matplotlib.pyplot as plt
+from yt.units import cm
+
+"""
+Given a plot file and field name, it gives slice plots at the top,
+middle, and bottom of the domain (shell).
+"""
+
+def slice(fname:str, field:str, loc:str="top") -> None:
+    """
+    A slice plot of the dataset for Spherical2D geometry.
+
+    Parameter
+    =======================
+    fname: plot file name
+    field: field parameter
+    loc: location on the domain. {top, mid, bot}
+    """
+
+    ds = yt.load(fname, hint="castro")
+    currentTime = ds.current_time.in_units("s")
+    print(f"Current time of this plot file is {currentTime} s")
+
+    # Some geometry properties
+    rr = ds.domain_right_edge[0].in_units("km")
+    rl = ds.domain_left_edge[0].in_units("km")
+    dr = rr - rl
+    r_center = 0.5 * (rr + rl)
+
+    thetar = ds.domain_right_edge[1]
+    thetal = ds.domain_left_edge[1]
+    dtheta = thetar - thetal
+    theta_center = 0.5 * (thetar + thetal)
+
+    # Domain width of the slice plot
+    width = (2.0*dr, 2.0*dr)
+
+    loc = loc.lower()
+    loc_options = ["top", "mid", "bot"]
+
+    if loc not in loc_options:
+        raise Exception("loc parameter must be top, mid or bot")
+
+    # Centers for the Top, Mid and Bot panels
+    centers = {"top":(r_center*np.sin(thetal)+0.8*dr, r_center*np.cos(thetal)),
+               "mid":(r_center*np.sin(theta_center), r_center*np.cos(theta_center)),
+               "bot":(r_center*np.sin(thetar)+0.8*dr, r_center*np.cos(thetar))}
+
+    sp = yt.SlicePlot(ds, 'phi', field, width=width)
+
+    sp.set_center(centers[loc])
+
+    sp.set_cmap(field, "viridis")
+    if field in ["x_velocity", "y_velocity", "z_velocity"]:
+        sp.set_cmap(field, "coolwarm")
+
+    # sp.annotate_text((0.05, 0.05), f"{currentTime.in_cgs():8.5f} s")
+    sp.save(f"{ds}_{loc}")
+
+if __name__ == "__main__":
+
+    if len(sys.argv) < 3:
+        raise Exception("Please enter parameters in order of: fname field_name loc[optional]")
+    fname = sys.argv[1]
+    field = sys.argv[2]
+    if len(sys.argv) > 3:
+        loc = sys.argv[3]
+
+    slice(fname, field, loc=loc)