#!/usr/bin/env python """Make a quick-look plot for a gamhelio run. Make a quick-look plot for a gamhelio run. Several different sets of plots are supported, and are distinguished by the value of the "pic" argument. pic1 (default): A 4-panel display showing pcolormesh plots in the z = 0 (equatorial) plane of the gamhelio frame used in the simulation. The plots are (r0 is the inner radius of the grid, which should be 21.5 Rsun): Upper left: Solar wind speed (km/s) Upper right: Solar wind number density scaled by (r/r0)**2 (cm**-3) Lower left: Solar wind temperature scaled by r/r0 (MK) Lower right: Solar wind radial magnetic field scaled by (r/r0)**2 (nT) pic2: A 4-panel display showing pcolormesh plots in the y = 0 (meridional, containing Earth) plane of the gamhelio frame used in the simulation. The plots are (r0 is the inner radius of the grid, which should be 21.5 Rsun): Upper left: Solar wind speed (km/s) Upper right: Solar wind number density scaled by (r/r0)**2 (cm**-3) Lower left: Solar wind temperature scaled by r/r0 (MK) Lower right: Solar wind radial magnetic field scaled by (r/r0)**2 (nT) pic3: A 4-panel display showing pcolormesh plots in the r = 1 AU slice of the gamhelio frame used in the simulation. The plots are: Upper left: Solar wind speed (km/s) Upper right: Solar wind number density (cm**-3) Lower left: Solar wind temperature (MK) Lower right: Solar wind radial magnetic field (nT) pic4: A pcolormesh plot in the innermost radial slice (r = 22 Rsun) of the gamhelio frame used in the simulation. The plot shows the radial magnetic field in nT, in a coordinate frame rotating with the Sun. - SCALED? pic5: A 3-panel display showing solar wind variables as a function of radius, 22 Rsun <= r <= 220 Rsun. The plots are: Upper left: Solar wind number density (cm**-3) - SCALED? Upper right: Solar wind speed (km/s) - SCALED? Lower left: Solar wind radial momentum flux (km**2/s**2/cm**3) - SCALED? pic6: A 4-panel display showing components of the solar wind magnetic field in the solar equatorial plane (z=0), for -200 Rsun <= X, Y <= +200 Rsun. Upper left: Radial component of magnetic field (nT) Upper right: x-component of magnetic field (nT) Lower left: y-component of magnetic field (nT) Lower right: z-component of magnetic field (nT) pic7: A 4-panel display showing pcolormesh plots in a j-slice. A j-slice is a slice through the gamhelio data cube at a fixed colatitude. j = 0 corresponds to the YZ plane of the gamhelio frame used in the simulation. The j = Nj/2-1 slice corresponds to the equatorial plane. The plots are: Upper left: Solar wind speed (km/s) Upper right: Solar wind number density scaled by (r/r0)**2 (cm**-3) Lower left: Solar wind temperature scaled by r/r0 (MK) Lower right: Solar wind radial magnetic field scaled by r/r0 (nT) Authors ------- Elena Provornikova (elena.provornikova@jhuapl.edu) Andrew McCubbin (andrew.mccubbin@jhuapl.edu) Eric Winter (eric.winter@jhuapl.edu) """ # Import standard modules. import argparse import os import time # Import supplemental modules. import astropy from astropy.coordinates import SkyCoord import astropy.units as u import matplotlib as mpl from matplotlib import gridspec import matplotlib.pyplot as plt import numpy as np import spacepy.datamodel as dm from sunpy.coordinates import frames # Import project-specific modules. from kaipy import cdaweb_utils import kaipy.gamhelio.helioViz as hviz import kaipy.gamhelio.heliosphere as hsph import kaipy.kaiH5 as kh5 import kaipy.kaiTools as ktools import kaipy.kaiViz as kv from kaipy.satcomp import scutils # Program constants and defaults # Program description. DESCRIPTION = "Make a quicklook plot for a gamhelio run." # Default identifier for results to read. DEFAULT_RUNID = "wsa" # List of steps DEFAULT_STEPS = "1" # Default slices DEFAULT_SLICE = None # Code for default picture type. DEFAULT_PICTYPE = "pic1" # Colors to use for spacecraft position symbols. SPACECRAFT_COLORS = list(mpl.colors.TABLEAU_COLORS.keys()) def create_command_line_parser(): """Create the command-line argument parser. Create the parser for command-line arguments. Parameters ---------- None Returns ------- parser : argparse.ArgumentParser Command-line argument parser for this script. Raises ------ None """ parser = argparse.ArgumentParser(description=DESCRIPTION) parser.add_argument( "--debug", action="store_true", help="Print debugging output (default: %(default)s)." ) parser.add_argument( "--directory", "-d", type=str, metavar="directory", default=os.getcwd(), help="Directory containing data to read (default: %(default)s)" ) parser.add_argument( "--hgsplot", action="store_true", help="Plot in the Heliographic Stonyhurst frame corresponding to the " "date of the plot (default: %(default)s)." ) parser.add_argument( "-id", type=str, metavar="runid", default=DEFAULT_RUNID, help="Run ID of data (default: %(default)s)" ) parser.add_argument( "-jslice", type=int, metavar="jSlice", default=None, help="Index of j-slice for pic7 (default: Nj/2-1)" ) parser.add_argument( "-lon", type=float, metavar="lon", default=0.0, help="Longitude of meridian slice (pic2) (default: %(default)s)" ) parser.add_argument( "--nlist", type=lambda n: [int(item) for item in n.split(',')], metavar="list of steps", default=DEFAULT_STEPS, help="List of time slice(s) n1,n2,... to plot (default: %(default)s)" ) parser.add_argument( "--nslice", type=lambda n: [int(item) for item in n.split(':')], metavar="step slice", default=DEFAULT_SLICE, help="Slice for range of time slice(s) n1:n2 to plot " "(default: %(default)s)" ) parser.add_argument( "--nworkers", "-nw", type=int, metavar="nworkers", default=4, help="Number of parallel workers (default: %(default)s)" ) parser.add_argument( "--parallel", "-p", action="store_true", help="Read from HDF5 in parallel (default: %(default)s)." ) parser.add_argument( "-pic", type=str, metavar="pictype", default=DEFAULT_PICTYPE, help="Code for plot type (pic1-pic7) (default: %(default)s)" ) parser.add_argument( "--spacecraft", type=str, metavar="spacecraft", default=None, help="Names of spacecraft to plot positions, separated by commas " "(default: %(default)s)" ) parser.add_argument( "--verbose", "-v", action="store_true", help="Print verbose output (default: %(default)s)." ) parser.add_argument( "--inner", action="store_true", help="Plot inner i-slice for pic3 (default: %(default)s)." ) # Return the parser. return parser def initFig(pic): """Create the matplotlib figure for a plot. Determine figure size (width, height) (inches) based on the pic type. Parameters ---------- pic : str String representing picture type. Returns ------- fig : mpl.Figure Matplotlib Figure to use for plots. Raises ------ None """ # Figure dimensions are in inches. if pic == "pic1" or pic == "pic2" or pic == "pic7": figSz = (10, 12.5) elif pic == "pic3": figSz = (10, 6.5) elif pic == "pic4": figSz = (10, 6) elif pic == "pic5": figSz = (12, 12) elif pic == "pic6": figSz = (10, 12.5) # Create the figure. fig = plt.figure(figsize=figSz, layout="constrained") return fig def fOut(runid, pic, nStp, hgsplot): """Compute the name of the output file. Compute the name of the output file. Parameters ---------- runid : str ID string for run pic : str String representing picture type. nStp : int Simulation step number used in plot. hgsplot : bool True if plot is in HGS frame at the date of the plot Returns ------- s : str Name of file to receive the plot. Raises ------ None """ if hgsplot: s = f"qkpic_{runid}_{pic}_HGS_n{nStp}.png" else: s = f"qkpic_{runid}_{pic}_n{nStp}.png" return s def GHtoHGS(mjd_gh, x_gh, y_gh, z_gh, mjd_hgs): """Convert Cartesin GH coordinates to HGS. Convert Cartesian coordinates in the gamhelio frame at time mjdc to the Heliographic Sonyhurst frame at time mjd. NOTE: The gamhelio frame at time t is related to the Heliographic Stonyhurst frame at time t by the reflection of the x- and y-axes: x_gh(t) = -x_hgs(t) y_gh(t) = -y_hgs(t) z_gh(t) = z_hgs(t) Since HGS is a time-dependent frame, a time must be provided for each set of coordinates. Parameters ---------- mjd_gh : float MJD of source gamhelio frame x_gh, y_gh, z_gh : np.array of float (any shape) or scalar float Cartesian coordinates in GH(mjdc) frame. All three arrays x, y, z must have identical shapes. mjd_hgs : float MJD of target HGS frame Returns ------- x_hgs, y_hgs, z_hgs : np.array of float (same shape as x_gh, y_gh, z_gh) Cartesian coordinates converted to HGS(mjd) frame. Raises ------ None """ # Load the source coordinates (originially in the GH(mjd_gh) frame) into # the equivalent HGS(mjd_gh) frame. c_gh = SkyCoord( -x_gh*u.Rsun, -y_gh*u.Rsun, z_gh*u.Rsun, frame=frames.HeliographicStonyhurst, obstime=ktools.MJD2UT(mjd_gh), representation_type="cartesian" ) # Create the target Heliographic Stonyhurst frame. hgs_frame = frames.HeliographicStonyhurst( obstime=ktools.MJD2UT(mjd_hgs) ) # Convert the coordinates from GH(mjd_gh) to HGS(mjd_hgs). c_hgs = c_gh.transform_to(hgs_frame) # Extract and return the converted coordinates. x_hgs = dm.dmarray(c_hgs.cartesian.x) y_hgs = dm.dmarray(c_hgs.cartesian.y) z_hgs = dm.dmarray(c_hgs.cartesian.z) return x_hgs, y_hgs, z_hgs def main(): """Make a quick-look plot for a gamhelio run.""" # Set up the command-line parser. parser = create_command_line_parser() # Parse the command-line arguments. args = parser.parse_args() if args.debug: print(f"args = {args}") debug = args.debug fdir = args.directory hgsplot = args.hgsplot ftag = args.id jslice = args.jslice pic2lon = args.lon steps = args.nlist slices = args.nslice nWorkers = args.nworkers doParallel = args.parallel pic = args.pic spacecraft = args.spacecraft verbose = args.verbose inner = args.inner if slices: print(f"Slice selected {slice(slices[0], slices[1], slices[2])}") # Fetch the plot domain based on the picture type. tic = time.perf_counter() xyBds = hviz.GetSizeBds(pic) toc = time.perf_counter() print(xyBds) print(f"Get bounds took {toc - tic} s") # Do work? doFast = False # Open a pipe to the results data. tic = time.perf_counter() gsph = hsph.GamsphPipe(fdir, ftag, doFast=doFast, doParallel=doParallel, nWorkers=nWorkers) toc = time.perf_counter() print(f"Open pipe took {toc-tic} s") # Compute the range of time steps to use. if slices and steps[0] == 1: steps = range(gsph.s0,gsph.sFin + 1)[slice(slices[0], slices[1], slices[2])] print(f"steps = {steps}") # Get the MJDc value for use in computing the gamhelio frame. fname = gsph.f0 MJDc = scutils.read_MJDc(fname) # Split the list into individual spacecraft names. if spacecraft: spacecraft = spacecraft.split(',') # Create figures in a memory buffer. mpl.use("Agg") # Make a plot for each time step in the list of time steps. for nStp in steps: if debug: print(f"nStp = {nStp}") tic = time.perf_counter() print(f"Generating {pic} for time step {nStp}.") fig = initFig(pic) # Extract the MJD for the step. if any(gsph.MJDs): mjd = gsph.MJDs[nStp-gsph.s0] time_stamp = ktools.MJD2UT(mjd) else: mjd = gsph.T[nStp-gsph.s0]/(3600./gsph.tScl) time_stamp = f"{mjd:0.2f} [hr]" if debug: print(f"mjd = {mjd}") # Lay out the subplots. if pic in ["pic1", "pic2", "pic3", "pic6", "pic7"]: gs = gridspec.GridSpec(4, 6, height_ratios=[20, 1, 20, 1], figure=fig) # Axes for plots. AxL0 = fig.add_subplot(gs[0, 0:3]) AxR0 = fig.add_subplot(gs[0, 3:]) AxL1 = fig.add_subplot(gs[2, 0:3]) AxR1 = fig.add_subplot(gs[2, 3:]) # Axes for colorbars. AxC1_0 = fig.add_subplot(gs[1, 0:3]) AxC2_0 = fig.add_subplot(gs[1, 3:]) AxC1_1 = fig.add_subplot(gs[3, 0:3]) AxC2_1 = fig.add_subplot(gs[3, 3:]) elif pic == "pic4": gs = gridspec.GridSpec(2, 1, height_ratios=[20, 1], figure=fig) Ax = fig.add_subplot(gs[0, 0]) AxC = fig.add_subplot(gs[1, 0]) elif pic == "pic5": gs = gridspec.GridSpec(2, 2, figure=fig) Ax = fig.add_subplot(gs[0, 0]) AxC = fig.add_subplot(gs[0, 1]) AxC1 = fig.add_subplot(gs[1, 0]) else: raise TypeError(f"Invalid figure type: {pic}!") # If the step is -1, use the last step. if nStp < 0: nStp = gsph.sFin print(f"Using Step {nStp}") # Now create the actual plots. if pic == "pic1": # Equatorial plots in the XY plane of the modified HGS frame used # by gamhelio. If hgsplot is True, then the plot frame is the true # HGS frame at the time of the plot. hviz.PlotEqMagV(gsph, nStp, xyBds, AxL0, AxC1_0, hgsplot=hgsplot, MJDc=MJDc, MJD_plot=mjd) hviz.PlotEqD(gsph, nStp, xyBds, AxR0, AxC2_0, hgsplot=hgsplot, MJDc=MJDc, MJD_plot=mjd) hviz.PlotEqTemp(gsph, nStp, xyBds, AxL1, AxC1_1, hgsplot=hgsplot, MJDc=MJDc, MJD_plot=mjd) hviz.PlotEqBr(gsph, nStp, xyBds, AxR1, AxC2_1, hgsplot=hgsplot, MJDc=MJDc, MJD_plot=mjd) if hgsplot: fig.suptitle("Heliographic Stonyhurst frame for " f"{time_stamp}") else: fig.suptitle(f"GAMERA-Helio frame for {time_stamp}") elif pic == "pic2": # Meridional plots in the XZ plane of the modified HGS frame used # by gamhelio. If hgsplot is True, then the plot frame is the true # HGS frame at the time of the plot. hviz.PlotMerMagV(gsph, nStp, xyBds, AxL0, AxC1_0, indx=(None, pic2lon), hgsplot=hgsplot, MJDc=MJDc, MJD_plot=mjd) hviz.PlotMerDNorm(gsph, nStp, xyBds, AxR0, AxC2_0, indx=(None, pic2lon), hgsplot=hgsplot, MJDc=MJDc, MJD_plot=mjd) hviz.PlotMerTemp(gsph, nStp, xyBds, AxL1, AxC1_1, indx=(None, pic2lon), hgsplot=hgsplot, MJDc=MJDc, MJD_plot=mjd) hviz.PlotMerBrNorm(gsph, nStp, xyBds, AxR1, AxC2_1, indx=(None, pic2lon), hgsplot=hgsplot, MJDc=MJDc, MJD_plot=mjd) if hgsplot: fig.suptitle("Heliographic Stonyhurst frame for " f"{time_stamp}") else: fig.suptitle(f"GAMERA-Helio frame for {time_stamp}") elif pic == "pic3": # Lat/lon plot at 1 AU (the outer edge of the gamhelio grid), in # the modified HGS frame rotating with the Sun. AU_RSUN = 215.0 radius = AU_RSUN I_RSUN = 21.5 if inner : radius = I_RSUN hviz.PlotiSlMagV(gsph, nStp, xyBds, AxL0, AxC1_0, idx=radius, idx_is_radius=True, hgsplot=hgsplot, MJDc=MJDc, MJD_plot=mjd) hviz.PlotiSlD(gsph, nStp, xyBds, AxR0, AxC2_0, idx=radius, idx_is_radius=True, hgsplot=hgsplot, MJDc=MJDc, MJD_plot=mjd, use_outer_range=(not inner) ) hviz.PlotiSlTemp(gsph, nStp, xyBds, AxL1, AxC1_1, idx=radius, idx_is_radius=True, hgsplot=hgsplot, MJDc=MJDc, MJD_plot=mjd, use_outer_range=(not inner)) hviz.PlotiSlBr(gsph, nStp, xyBds, AxR1, AxC2_1, idx=radius, idx_is_radius=True, hgsplot=hgsplot, MJDc=MJDc, MJD_plot=mjd, use_outer_range=(not inner)) if hgsplot: fig.suptitle(f"Heliographic Stonyhurst frame at {radius} [RE] for {time_stamp}") else: fig.suptitle(f"GAMERA-Helio frame at {radius} [RE] for {time_stamp}") elif pic == "pic4": # Plot at 1 AU in frame rotating with Sun. hviz.PlotiSlBrRotatingFrame(gsph, nStp, xyBds, Ax, AxC) elif pic == "pic5": hviz.PlotDensityProf(gsph, nStp, xyBds, Ax) hviz.PlotSpeedProf(gsph, nStp, xyBds, AxC) hviz.PlotFluxProf(gsph, nStp, xyBds, AxC1) elif pic == "pic6": hviz.PlotEqBr(gsph, nStp, xyBds, AxL0, AxC1_0, hgsplot=hgsplot, MJDc=MJDc, MJD_plot=mjd) hviz.PlotEqBx(gsph, nStp, xyBds, AxR0, AxC2_0, hgsplot=hgsplot, MJDc=MJDc, MJD_plot=mjd) hviz.PlotEqBy(gsph, nStp, xyBds, AxL1, AxC1_1, hgsplot=hgsplot, MJDc=MJDc, MJD_plot=mjd) hviz.PlotEqBz(gsph, nStp, xyBds, AxR1, AxC2_1, hgsplot=hgsplot, MJDc=MJDc, MJD_plot=mjd) fig.suptitle("GAMERA-Helio frame for " f"{time_stamp}") elif pic == "pic7": if jslice is None: jidx = gsph.Nj//2 - 1 else: jidx = jslice hviz.PlotjMagV(gsph, nStp, xyBds, AxL0, AxC1_0, jidx=jidx, hgsplot=hgsplot, MJDc=MJDc, MJD_plot=mjd) hviz.PlotjD(gsph, nStp, xyBds, AxR0, AxC2_0, jidx=jidx, hgsplot=hgsplot, MJDc=MJDc, MJD_plot=mjd) hviz.PlotjTemp(gsph, nStp, xyBds, AxL1, AxC1_1, jidx=jidx, hgsplot=hgsplot, MJDc=MJDc, MJD_plot=mjd) hviz.PlotjBr(gsph, nStp, xyBds, AxR1, AxC2_1, jidx=jidx, hgsplot=hgsplot, MJDc=MJDc, MJD_plot=mjd) fig.suptitle("GAMERA-Helio frame for " f"{time_stamp}") else: raise TypeError(f"Invalid figure type: {pic}!") # Add time in the upper left (if not in figure title). # if pic == "pic1" or pic == "pic2" or pic == "pic3" or pic == "pic6": if pic == "pic4" or pic == "pic5": gsph.AddTime(nStp, Ax, xy=[0.015, 0.92], fs="small") # Overlay the spacecraft positions. if spacecraft: # Fetch the MJD at start and end of the model results. MJD_start = kh5.tStep(fname, gsph.s0, aID="MJD") MJD_end = kh5.tStep(fname, gsph.sFin, aID="MJD") # Convert the start and stop MJD to a datetime object in UT. ut_start = ktools.MJD2UT(MJD_start) ut_end = ktools.MJD2UT(MJD_end) # Fetch the trajectory of each spacecraft from CDAWeb. Then # interpolate the position at the time of the plot, and plot the # spacecraft at the interpolated position. for (i_sc, sc_id) in enumerate(spacecraft): if verbose: print(f"Fetching trajectory for {sc_id}.") # Fetch the spacecraft trajectory in whatever frame is # available from CDAWeb. sc_data = cdaweb_utils.fetch_helio_spacecraft_trajectory( sc_id, ut_start, ut_end ) if sc_data is None: print(f"No trajectory found for {sc_id}.") continue # Ingest the trajectory by converting it to the GH(MJDc) frame. if verbose: print(f"Converting ephemeris for {sc_id} to the gamhelio " f"frame at MJD {MJDc}.") t_strings = np.array([str(t) for t in sc_data["Epoch"]]) t = astropy.time.Time(t_strings, scale='utc').mjd x, y, z = cdaweb_utils.ingest_helio_spacecraft_trajectory( sc_id, sc_data, MJDc) # Interpolate the spacecraft position at the time for the plot. t_sc = mjd x_sc = np.interp(t_sc, t, x) y_sc = np.interp(t_sc, t, y) z_sc = np.interp(t_sc, t, z) # If needed, convert the position to HGS(mjd). if hgsplot: x_sc, y_sc, z_sc = GHtoHGS(MJDc, x_sc, y_sc, z_sc, mjd) # If needed, compute heliocentric spherical coordinates # for the interpolated spacecraft position. Longitude is in # the -180 to +180 range. Convert to 0-360 if not using # hgsplot. if pic == "pic3" or pic == "pic4": rxy = np.sqrt(x_sc**2 + y_sc**2) theta = np.arctan2(rxy, z_sc) phi = np.arctan2(y_sc, x_sc) lat = np.degrees(np.pi/2 - theta) lon = np.degrees(phi) lat_sc = lat lon_sc = lon if not hgsplot: if lon_sc < 0: lon_sc += 360 # Plot the position of the spacecraft at the plot time. Each # spacecraft is plotted as a colored dot with a black outline. color = SPACECRAFT_COLORS[i_sc % len(SPACECRAFT_COLORS)] x_nudge = 0.0 y_nudge = 8.0 if pic == "pic1": for ax in (AxL0, AxR0, AxL1, AxR1): ax.plot(x_sc, y_sc, 'o', c=color) ax.plot(x_sc, y_sc, 'o', c="black", fillstyle="none") ax.text(x_sc + x_nudge, y_sc + y_nudge, sc_id, c="black", horizontalalignment="center") elif pic == "pic2": for ax in (AxL0, AxR0, AxL1, AxR1): ax.plot(x_sc, z_sc, 'o', c=color) ax.plot(x_sc, z_sc, 'o', c="black", fillstyle="none") ax.text(x_sc + x_nudge, z_sc + y_nudge, sc_id, c="black", horizontalalignment="center") elif pic == "pic3": for ax in (AxL0, AxR0, AxL1, AxR1): ax.plot(lon_sc, lat_sc, 'o', c=color) ax.plot(lon_sc, lat_sc, 'o', c="black", fillstyle="none") ax.text(lon_sc + x_nudge, lat_sc + y_nudge, sc_id, c="black", horizontalalignment="center") elif pic == "pic4": ax = Ax ax.plot(lon_sc, lat_sc, 'o', c=color) ax.plot(lon_sc, lat_sc, 'o', c="black", fillstyle="none") ax.text(lon_sc + x_nudge, lat_sc + y_nudge, sc_id, c="black", horizontalalignment="center") elif pic == "pic5": pass elif pic == "pic6": for ax in (AxL0, AxR0, AxL1, AxR1): ax.plot(x_sc, y_sc, 'o', c=color) ax.plot(x_sc, y_sc, 'o', c="black", fillstyle="none") ax.text(x_sc + x_nudge, y_sc + y_nudge, sc_id, c="black", horizontalalignment="center") elif pic == "pic7": raise TypeError("Spacecraft not supported for pic7!") else: raise TypeError(f"Invalid plot code: {pic}!") # Save the figure to a file. path = os.path.join(fdir, fOut(ftag, pic, nStp, hgsplot)) kv.savePic(path, bLenX=40) plt.close() toc = time.perf_counter() print(f"Step {nStp} took {toc-tic} s") if __name__ == "__main__": main()