mirror of
https://github.com/JHUAPL/kaiju.git
synced 2026-01-08 22:58:05 -05:00
658 lines
24 KiB
Python
Executable File
658 lines
24 KiB
Python
Executable File
#!/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()
|