def adjust_plots(ax, ylab, xlab=False, Zero=True):
ax.grid(True)
applySmartTimeTicks(ax, timerange)
ax.set_ylabel(ylab)
labels = ax.get_yticklabels()
labels[-1].set_visible(False)
labels[0].set_visible(False)
if Zero:
ax.plot(timerange, [0, 0], 'k--')
if xlab:
ax.set_xlabel('Universal Time from %s' % timerange[0].isoformat())
else:
ax.xaxis.set_ticklabels([])
def main(infiles=None):
# read SWMF ImfInput file
# Originally written to examine data from simulations
# by Morley, Welling and Woodroffe (2018). See data at
# Zenodo (https://doi.org/10.5281/zenodo.1324562)
infilename = 'Event5Ensembles/run_orig/IMF.dat'
eventIMF = bats.ImfInput(filename=infilename)
data1 = bats.LogFile('Event5Ensembles/run_orig/GM/IO2/log_e20100404-190000.log')
# #read IMF for 10 events...
# infiles = ['Event5Ensembles/run_{:03d}/IMF.dat'.format(n)
# for n in [32,4,36,10,13,17,18,20,24,29]]
# read IMF files for all ensemble members
if infiles is None:
infiles = glob.glob('Event5Ensembles/run_???/IMF.dat')
subsetlabel = False
else:
# got list of run directories from Dst/Kp plotter
subsetlabel = True
infiles = [os.path.join(d, 'IMF.dat') for d in infiles]
nsubset = len(infiles)
eventlist = [bats.ImfInput(filename=inf) for inf in infiles]
tstart = eventIMF['time'][0]
tstop = eventIMF['time'][-1]
sym = kyo.KyotoSym(lines=bats.kyoto.symfetch(tstart, tstop))
fig = plt.figure(figsize=(10, 5))
ax1 = fig.add_subplot(211)
ax2 = fig.add_subplot(212)
for ev in eventlist:
gco = '{}'.format(np.random.randint(5, 50)/100.0)
pred01B = Dst_Burton(sym['sym-h'][0], ev['ux'], ev['bz'], dt=1./60)
pred01O = Dst_OBrien(sym['sym-h'][0], ev['ux'], ev['bz'], dt=1./60)
ax1.plot(ev['time'], pred01B, c=gco, alpha=0.5)
ax2.plot(ev['time'], pred01O, c=gco, alpha=0.5)
# ax1.plot(sym['time'], sym['sym-h'], lw=1.5, c='crimson', label='Sym-H')
evtime = spt.Ticktock(eventIMF['time']).RDT
datime = spt.Ticktock(data1['time']).RDT
simDst = tb.interpol(evtime, datime, data1['dst'])
# ax1.plot(eventIMF['time'], simDst+11-(7.26*eventIMF['pram']), lw=1.5, c='seagreen', label='Sym-H (Press.Corr.)')
# ax2.plot(sym['time'], sym['sym-h'], lw=1.5, c='crimson')
# ax2.plot(eventIMF['time'], simDst+11-(7.26*eventIMF['pram']), lw=1.5, c='seagreen', label='Sym-H (Press.Corr.)')
ax1.plot(data1['time'], data1['dst'], linewidth=1.5, color='crimson', alpha=0.65, label='SWMF')
ax2.plot(data1['time'], data1['dst'], linewidth=1.5, color='crimson', alpha=0.65)
ax1.legend()
splot.applySmartTimeTicks(ax1, [tstart, tstop])
splot.applySmartTimeTicks(ax2, [tstart, tstop], dolabel=True)
ax1.set_ylabel('Sym-H [nT]')
ax2.set_ylabel('Sym-H [nT]')
ax1.text(0.05, 0.05, "Burton et al.", transform=ax1.transAxes)
ax2.text(0.05, 0.05, "O'Brien et al.", transform=ax2.transAxes)
def plotEvent(st=dt.datetime(1989, 3, 11), en=dt.datetime(1989, 3, 15)):
datapath = os.path.abspath(os.path.join('..', 'ref_data', '1989'))
# When plotting, use spacepy.datamanager to insert fill between contiguous regions
fig, ax = plt.subplots(5, sharex=True, figsize=(12, 8))
# Do MIT plasma
mit_pl = readIMP8plasmafile(
os.path.join(datapath, 'imp8.data.1989.060.090'))
mit_t, mit_sp = dman.insert_fill(np.asarray(mit_pl['time']),
np.asarray(mit_pl['speed']))
ax[0].plot(mit_t, mit_sp, label='MIT plasma')
mit_t, mit_nd = dman.insert_fill(np.asarray(mit_pl['time']),
np.asarray(mit_pl['n_dens']))
ax[1].plot(mit_t, mit_nd, label='MIT plasma')
# Do LANL plasma
lanl_pl = readIMPplasmaLANL(
os.path.join(datapath, '198903_imp8_lanl_sw_2min.asc'))
lanl_t, lanl_sp = dman.insert_fill(np.asarray(lanl_pl['time']),
np.asarray(lanl_pl['speed']))
ax[0].plot(lanl_t, lanl_sp, label='LANL plasma')
lanl_t, lanl_nd = dman.insert_fill(np.asarray(lanl_pl['time']),
np.asarray(lanl_pl['n_dens']))
ax[1].plot(lanl_t, lanl_nd, label='LANL plasma')
# Do IMF from ISEE-3
isee_ma = readISEEmag(os.path.join(datapath,
'198903_isee3_mag03_1min.asc'))
ax[2].plot(isee_ma['time'], isee_ma['B'][:, 0], label='ISEE-3 Bx')
ax[2].plot(isee_ma['time'], isee_ma['B'][:, 1], label='ISEE-3 By')
ax[2].plot(isee_ma['time'], isee_ma['B'][:, 2], label='ISEE-3 Bz')
ax[2].axhline(linestyle='--', color=(0.3, 0.3, 0.3))
# Add stuff from burton_test
from burton import invert_example
invert_example(axes=[ax[3], ax[4]], show=False)
# Finalize
ax[0].legend(loc='upper left', fancybox=True, framealpha=0.5)
ax[0].set_ylabel('Speed\n[km/s]')
ax[1].legend(loc='upper left', fancybox=True, framealpha=0.5)
ax[1].set_ylabel('Number Density\n[cm$^{-3}$]')
ax[2].legend(loc='upper left', fancybox=True, framealpha=0.5)
ax[2].set_ylabel('ISEE3\nIMF [nT]')
ax[2].set_ylim([-15, 15])
splot.applySmartTimeTicks(ax[2], [st, en], dolabel=False)
ax[2].set_xlim([st, en])
plt.show()
def comp_mag(name, obs, mod, interactive=False):
'''
Given a magnetometer with station name "name", compare the
data and model together and save the plot.
'''
import matplotlib.pyplot as plt
# Create a figure with 3 subplots
fig = plt.figure(figsize=(10, 10))
a1, a2, a3 = fig.subplots(3, 1)
# Loop over field component; plot data v. model for each
for x1, x2, ax in zip('ned', 'xyz', (a1, a2, a3)):
# Plot model & data:
ax.plot(mod['time'], mod['dB'+x1], label='SWMF')
ax.plot(obs['time'], obs[name]['b'+x2], label='SuperMag')
ax.set_ylabel(f'$\\Delta B_{x2}$ ($nT$)')
# Adjust
applySmartTimeTicks(ax, mod['time'], dolabel=ax == a3)
# Add legend/title
a1.legend(loc='best')
# Build title for our plot using mag_info (if available):
if name in mag_info:
info = mag_info[name]
title = f"{info['station-name']} ({name}) " + \
f"(mlat={info['aacgmlat']:.1f}$^{{\\circ}}$, " + \
f" mlon={info['aacgmlon']:.1f}$^{{\\circ}}$)"
else:
title = f"Station {name}"
a1.set_title(title)
fig.tight_layout()
if not interactive:
fig.savefig(f'{args.outdir}/{name}.png')
plt.close('all')
stations = ['YKC', 'MEA', 'NEW', 'FRN', 'IQA', 'PBQ', 'OTT', 'FRD', 'VIC']
stations = [key for key in origdata.keys() if len(key)==3]
for stat in stations:
if stat not in smdata.station: continue
subset = origdata[stat]
simtime = subset['time'][::6]
dBdte = decimate(subset['dBdte'], 6)
dBdtn = decimate(subset['dBdtn'], 6)
dBdth = np.array([linalg.norm([dBdtn[i], dBdte[i]]) for i in range(len(dBdtn))])
smstat = smdata.station[stat]
Bdoth = np.array([linalg.norm(smstat['Bdot'][i,:2]) for i in range(len(smstat['Bdot']))])
fig = plt.figure(figsize=(10,4))
ax = fig.add_subplot(111)
ax.plot(simtime, dBdth, 'b-', alpha=0.4)
ax.plot(smstat['time'], Bdoth, 'r-', alpha=0.4)
run20, t20 = tb.windowMean(dBdth, time=simtime, winsize=dt.timedelta(minutes=20), overlap=dt.timedelta(0), st_time=dt.datetime(2000,8,12,1), op=np.max)
ax.plot(t20, run20, marker='o', color='b', linestyle='none', markersize=3, label='SWMF')
obs20, t20 = tb.windowMean(Bdoth, time=smstat['time'], winsize=dt.timedelta(minutes=20), overlap=dt.timedelta(0), st_time=dt.datetime(2000,8,12,1), op=np.max)
ax.plot(t20, obs20, marker='x', color='r', linestyle='none', markersize=3, label='Obs')
ax.set_ylabel('1-min dB/dt$_{H}$ [nT/s]')
ax.set_xlabel('2000-08-12')
splot.applySmartTimeTicks(ax, subset['time'], dolimit=True)
plt.legend()
plt.title(stat)
plt.tight_layout()
#plt.show()
plt.savefig('Aug2000_dBdth_{}.png'.format(stat))
plt.close()
def add_ltut(self,
target=None,
loc=111,
cmap='Greens_r',
zlim=[1, 1000],
add_cbar=True,
clabel='Density $cm^{-3}$',
xlabel='full',
ylim=[4, 20],
title=None,
grid=True,
ntick=5):
'''
Plot log(density) as a contour against local time (y-axis) and
universal time (x-axis) using the PyBats *target* method of other
standard plotting methods. Four items are returned: the Matplotlib
Figure, Axes, Mesh, and ColorBar objects used (if cbar is set to
**False**, the returned ColorBar object is simply set to **False**.)
========== =======================================================
Kwarg Description
---------- -------------------------------------------------------
target Select plot destination. Defaults to new figure/axis.
loc The location of any generated subplots. Default is 111.
add_cbar Toggles the automatic colorbar. Default is**True**.
cmap Selects Matplotlib color table. Defaults to *Greens_r*.
zlim Limits for z-axis. Defaults to [0.1, 1000]
ylim Sets the MLT range on the y-axis. Defaults to [4,20].
clabel Sets colorbar label. Defaults to units.
xlabel Sets x-axis labels, use 'full', 'ticks', or **None**.
title Sets axis title; defaults to **None**.
grid Show white dotted grid? Defaults to **True**
ntick Number of attempted cbar ticks. Defaults to 5.
========== =======================================================
'''
import matplotlib.pyplot as plt
from matplotlib.ticker import LogLocator, LogFormatterMathtext
from matplotlib.colors import LogNorm
# Set ax and fig based on given target.
fig, ax = set_target(target, loc=loc)
# Enforce values to be within limits.
z = np.where(self['n'] > zlim[0], self['n'], 1.01 * zlim[0])
z[z > zlim[1]] = zlim[1]
# Create plot:
mesh = ax.pcolormesh(self._dtime,
self._y,
z.transpose(),
cmap=plt.get_cmap(cmap),
norm=LogNorm(),
vmin=zlim[0],
vmax=zlim[-1])
# Use LT ticks and markers on y-axis:
ax.set_yticks([6, 12, 18])
ax.set_yticklabels(['Dawn', 'Noon', 'Dusk'])
ax.set_ylim(ylim)
# White ticks, slightly thicker:
ax.tick_params(axis='both', which='both', color='w', width=1.2)
# Grid marks:
if grid: ax.grid(c='w')
if title: ax.set_title(title)
if xlabel == 'full':
# Both ticks and label.
applySmartTimeTicks(ax, self['time'], dolabel=True)
elif xlabel == 'ticks':
# Ticks, but no date label.
applySmartTimeTicks(ax, self['time'], dolabel=False)
else:
# A blank x-axis is often useful.
applySmartTimeTicks(ax, self['time'], dolabel=False)
ax.set_xticklabels('')
# Add cbar as necessary:
if add_cbar:
#lct = LogLocator
cbar = plt.colorbar(mesh, ax=ax, pad=0.01,
shrink=0.85) #, ticks=lct)
cbar.set_label(clabel)
else:
cbar = None
return fig, ax, mesh, cbar
def add_histplot(self,
target=False,
loc=111,
label='Kyoto $K_{p}$',
time_range=None,
filled=False,
level_kwargs={},
**kwargs):
'''
Make a quick histogram-style plot of the Kp data.
Returns
=======
fig : matplotlib figure object
ax : matplotlib axes object
Other Parameters
================
target : Figure or Axes
If None (default), a new figure is generated from scratch.
If a matplotlib Figure object, a new axis is created
to fill that figure.
If a matplotlib Axes object, the plot is placed
into that axis.
loc : int
Use to specify the subplot placement of the axis
(e.g. loc=212, etc.) Used if target is a Figure or None.
Default 111 (single plot).
label : string
The label applied to the line when a legend is added to the axes.
Defaults to 'Kyoto $K_{p}$'.
time_range : tuple of datetimes
The time range to plot. Only the first and last values in the
tuple (or list) are used if the number of elements is greater than
two. Defaults to **None**, meaning that the full time range
available is used.
filled : Boolean
If True, make a filled 'traffic light' plot of Kp using
spacepy.plot.levelPlot. Extra keyword arguments for levelPlot can
be supplied via level_kwargs.
Extra keyword arguments are passed to :function:`matplotlib.pyplot.plot`
to customize the line style (and are ignored if filled is True).
Examples
========
>>> import matplotlib.pyplot as plt
>>> import spacepy.pybats.kyoto as kt
>>> kp = kt.fetch('kp', (1981, 11), (1981, 11))
>>> kp.add_histplot(lw=2.0, color='r', label='Example Kp')
>>> ax = plt.gca()
>>> kp.add_histplot(filled=True)
'''
import matplotlib.pyplot as plt
if not time_range:
time_range = self['time']
fig, ax = set_target(target, figsize=(10, 4), loc=loc)
if not filled:
tvar = self['binstart'].tolist()
tvar.append(tvar[-1] + dt.timedelta(hours=3))
kpvar = self['kp'].tolist()
kpvar.append(kpvar[-1])
line = ax.plot(tvar,
kpvar,
label=label,
drawstyle='steps-post',
**kwargs)
else:
ax = levelPlot(self,
time='binstart',
var='kp',
target=ax,
**level_kwargs)
applySmartTimeTicks(ax, time_range, dolabel=True)
return fig, ax
def plotSummary(self,
timerange=None,
coord_sys=None,
fig_target=None,
spec=False,
orbit_params=(False, True),
**kwargs):
"""Generate summary plot of AE9/AP9/SPM data loaded
spec : if True, plot spectrogram instead of flux/fluence lineplot, requires 'ecol' keyword
"""
if timerange:
if isinstance(timerange, spt.Ticktock):
t1 = timerange.UTC[0]
t2 = timerange.UTC[-1]
elif isinstance(timerange[0], dt.datetime):
t1 = timerange[0]
t2 = timerange[-1]
else:
raise TypeError('Incorrect data type provided for timerange')
# now select subset
i_use = tb.tOverlapHalf([t1, t2], self['Epoch'])
t_use = self['Epoch'][i_use]
c_use = self['Coords'][i_use]
f_use = self[self.attrs['varname']][i_use, ...]
else:
t_use = self['Epoch']
c_use = self['Coords']
f_use = self[self.attrs['varname']]
if coord_sys and (coord_sys.upper() != c_use.attrs['COORD_SYS']):
# TODO: We assume cartesian, make flexible so can take spherical
cx = spc.Coords(c_use, c_use.attrs['COORD_SYS'], 'car')
cx.ticks = spt.Ticktock(t_use)
cx = cx.convert(coord_sys.upper(), 'car').data
else:
coord_sys = c_use.attrs['COORD_SYS']
cx = c_use
sys_subs = r'$_{' + coord_sys + r'}$'
splot.style('spacepy')
if orbit_params[0]:
landscape = orbit_params[1]
locs = [121, 122] if landscape else [211, 212]
else:
locs = [223, 224]
landscape = True
if fig_target:
fig, ax1 = splot.set_target(fig_target, loc=locs[0])
else:
fig, ax1 = splot.set_target(fig_target,
figsize=(8, 8),
loc=locs[0])
fig, ax2 = splot.set_target(fig, loc=locs[1])
ax1 = self._makeOrbitAxis(cx[:, 0], cx[:, 1], ax1)
ax2 = self._makeOrbitAxis(cx[:, 0], cx[:, 2], ax2)
if landscape:
ax1.set_xlabel('X{0} [{1}]'.format(sys_subs,
self['Coords'].attrs['UNITS']))
ax2.set_xlabel('X{0} [{1}]'.format(sys_subs,
self['Coords'].attrs['UNITS']))
ax1.set_ylabel('Y{0} [{1}]'.format(sys_subs,
self['Coords'].attrs['UNITS']))
ax2.set_ylabel('Z{0} [{1}]'.format(sys_subs,
self['Coords'].attrs['UNITS']))
ax1xl, ax1yl, ax2xl, ax2yl = ax1.get_xlim(), ax1.get_ylim(
), ax2.get_xlim(), ax2.get_ylim()
maxabslim = np.max(np.abs([ax1xl, ax1yl, ax2xl, ax2yl]))
if np.abs((ax1.get_xlim()[1] - ax1.get_xlim()[0])) < 1:
refpt = maxabslim
ax1.set_xlim([-1.25 * refpt, 1.25 * refpt])
ax1.set_ylim(ax1.get_xlim())
refpt = ax2.get_xlim()[0]
ax2.set_xlim([-1.25 * refpt, 1.25 * refpt])
ax2.set_ylim(ax1.get_xlim())
else:
ax1.set_xlim([-maxabslim, maxabslim])
ax1.set_ylim([-maxabslim, maxabslim])
ax2.set_xlim([-maxabslim, maxabslim])
ax2.set_ylim(ax2.get_xlim())
ax1.invert_yaxis()
ax1.invert_xaxis()
ax2.invert_xaxis()
ax1.set_aspect('equal')
ax2.set_aspect('equal')
if not orbit_params[0]:
ax3 = splot.plt.subplot2grid((2, 2), (0, 0), colspan=2)
if not spec:
l3 = ax3.semilogy(t_use, f_use)
ylab = '{0} ['.format(self.attrs['varname']) + re.sub(
'(\^[\d|-]*)+', _grp2mathmode,
self[self.attrs['varname']].attrs['UNITS']) + ']'
ax3.set_ylabel(ylab)
for ll, nn in zip(l3, self['Energy']):
ll.set_label('{0} {1}'.format(
nn, self['Energy'].attrs['UNITS']))
ncol = len(self['Energy']) // 2 if len(
self['Energy']) <= 6 else 3
leg = ax3.legend(loc='center',
bbox_to_anchor=(0.5, 1),
ncol=ncol,
frameon=True,
framealpha=0.5)
lims3 = ax3.get_ylim()
newupper = 10**(np.log10(lims3[0]) +
(np.log10(lims3[1] / lims3[0]) * 1.125))
ax3.set_ylim([lims3[0], newupper])
splot.applySmartTimeTicks(ax3, t_use)
fig.tight_layout()
pos3 = ax3.get_position()
ax3.set_position(
[pos3.x0, pos3.y0, pos3.width, pos3.height * 0.8])
# fig.suptitle('{model_type}\n'.format(**self.attrs) +
# '{0} - {1}'.format(t_use[0].isoformat()[:19], t_use[-1].isoformat()[:19]))
else:
if timerange:
raise NotImplementedError(
'Time range selection not yet implemented for spectrograms'
)
pos3 = ax3.get_position()
ax3.set_position(
[pos3.x0, pos3.y0, pos3.width, pos3.height * 0.9])
ecol = kwargs['ecol'] if 'ecol' in kwargs else 0
ax3 = self.plotSpectrogram(target=ax3, ecol=ecol)
splot.plt.subplots_adjust(wspace=0.3)
pos1 = ax1.get_position()
ax1.set_position([pos3.x0, pos1.y0, pos1.width, pos1.height])
splot.revert_style()
return fig
def plot_mags():
'''
Plot all magnetometers and save to file.
'''
path = './mag_summaries/'
# Create output directory as needed:
if not os.path.exists(path):
os.mkdir(path)
# Open magnetometer files:
mag_path = glob('GM/magnetometers*.mag')
if not mag_path:
raise ValueError('Could not find magnetometers in run directory.')
mags = bats.MagFile(mag_path[0])
mags.calc_h()
mags.calc_dbdt()
# Open observations as necessary:
do_obs = bool(args.supermag)
if do_obs:
data = read_supermag(args.supermag)
# Loop through magnetometers:
stations = mags.attrs['namemag']
for s in stations:
# shortcut!
mag = mags[s]
# Create figure & axes:
fig = plt.figure(figsize=(10, 7.5))
fig.subplots_adjust(top=0.931,
bottom=0.084,
left=0.101,
right=0.955,
hspace=0.141,
wspace=0.2)
a1, a2 = fig.subplots(2, 1)
# Plot models:
a1.plot(mag['time'], mag['dBh'], lw=2)
a2.plot(mag['time'], mag['dBdth'], lw=2)
# Plot observations:
if do_obs:
if s in data:
a1.plot(data['time'], data[s + '_H'], 'k', alpha=.7)
a2.plot(data['time'], data[s + '_dH'], 'k', alpha=.7)
# Customize axes:
applySmartTimeTicks(a1, mag['time'], dolabel=False)
applySmartTimeTicks(a2, mag['time'], dolabel=True)
a1.set_ylabel(r'$\Delta B_H$ ($nT$)', size=20)
a2.set_ylabel(r'$|dB/dt|_H$ ($nT/s$)', size=20)
# Add title:
a1.set_title('Station {0} on {1:%Y-%m-%d}'.format(s, mag['time'][0]),
size=20)
# Add legend:
if do_obs:
a1.legend(['Observations', 'SWMF'], loc='best')
# Save figure:
fig.savefig(path + '{}.png'.format(s))
plt.close('all')
def add_ltut(self, target=None, loc=111, cmap='Greens_r', zlim=[1,1000],
add_cbar=True, clabel='Density $cm^{-3}$', xlabel='full',
ylim=[4,20], title=None, grid=True, ntick=5):
'''
Plot log(density) as a contour against local time (y-axis) and
universal time (x-axis) using the PyBats *target* method of other
standard plotting methods. Four items are returned: the Matplotlib
Figure, Axes, Mesh, and ColorBar objects used (if cbar is set to
**False**, the returned ColorBar object is simply set to **False**.)
========== =======================================================
Kwarg Description
---------- -------------------------------------------------------
target Select plot destination. Defaults to new figure/axis.
loc The location of any generated subplots. Default is 111.
add_cbar Toggles the automatic colorbar. Default is**True**.
cmap Selects Matplotlib color table. Defaults to *Greens_r*.
zlim Limits for z-axis. Defaults to [0.1, 1000]
ylim Sets the MLT range on the y-axis. Defaults to [4,20].
clabel Sets colorbar label. Defaults to units.
xlabel Sets x-axis labels, use 'full', 'ticks', or **None**.
title Sets axis title; defaults to **None**.
grid Show white dotted grid? Defaults to **True**
ntick Number of attempted cbar ticks. Defaults to 5.
========== =======================================================
'''
import matplotlib.pyplot as plt
from matplotlib.ticker import LogLocator, LogFormatterMathtext
from matplotlib.colors import LogNorm
# Set ax and fig based on given target.
fig, ax = set_target(target, loc=loc)
# Enforce values to be within limits.
z=np.where(self['n']>zlim[0], self['n'], 1.01*zlim[0])
z[z>zlim[1]] = zlim[1]
# Create plot:
mesh = ax.pcolormesh(self._dtime, self._y, z.transpose(),
cmap=plt.get_cmap(cmap), norm=LogNorm(),
vmin=zlim[0], vmax=zlim[-1])
# Use LT ticks and markers on y-axis:
ax.set_yticks([6, 12, 18])
ax.set_yticklabels(['Dawn', 'Noon', 'Dusk'])
ax.set_ylim(ylim)
# White ticks, slightly thicker:
ax.tick_params(axis='both', which='both', color='w', width=1.2)
# Grid marks:
if grid: ax.grid(c='w')
if title: ax.set_title(title)
if xlabel == 'full':
# Both ticks and label.
applySmartTimeTicks(ax, self['time'], dolabel=True)
elif xlabel == 'ticks':
# Ticks, but no date label.
applySmartTimeTicks(ax, self['time'], dolabel=False)
else:
# A blank x-axis is often useful.
applySmartTimeTicks(ax, self['time'], dolabel=False)
ax.set_xticklabels('')
# Add cbar as necessary:
if add_cbar:
#lct = LogLocator
cbar=plt.colorbar(mesh, ax=ax, pad=0.01, shrink=0.85)#, ticks=lct)
cbar.set_label(clabel)
else:
cbar=None
return fig, ax, mesh, cbar
def add_histplot(self, target=False, loc=111, label='Kyoto $K_{p}$',
time_range=None, **kwargs):
'''
Make a quick histogram-style plot of the Kp data.
Returns
=======
fig : matplotlib figure object
ax : matplotlib axes object
Other Parameters
================
target : Figure or Axes
If None (default), a new figure is generated from scratch.
If a matplotlib Figure object, a new axis is created
to fill that figure.
If a matplotlib Axes object, the plot is placed
into that axis.
loc : int
Use to specify the subplot placement of the axis
(e.g. loc=212, etc.) Used if target is a Figure or None.
Default 111 (single plot).
label : string
The label applied to the line when a legend is added to the axes.
Defaults to 'Kyoto $K_{p}$'.
time_range : tuple of datetimes
The time range to plot. Only the first and last values in the
tuple (or list) are used if the number of elements is greater than
two. Defaults to **None**, meaning that the full time range
available is used.
Extra keyword arguments are passed to :function:`matplotlib.pyplot.plot`
to customize the line style.
Examples
========
>>> import spacepy.pybats.kyoto as kt
>>> kp = kt.fetch('kp', (1981, 11), (1981, 11)
>>> kp.add_histplot(lw=2.0, lc='r', label='Example KP')
'''
import matplotlib.pyplot as plt
# Shortcuts for the lazy.
bstart=self['binstart']
bstop =self['binstop']
npts =self.attrs['npts']
# Reformulate time to get histogram-type look.
newtime=np.zeros(npts*24, dtype=object)
newkp =np.zeros(npts*24)
# Set time range.
if not time_range:
time_range = newtime
for i in range(npts*8):
newtime[3*i ] = bstart[i]
newtime[3*i+1] = self['time'][i]
newtime[3*i+2] = bstop[i]
newkp[3*i:3*i+3] = self['kp'][i], self['kp'][i], self['kp'][i]
fig, ax = set_target(target, figsize=(10,4), loc=loc)
line=ax.plot(newtime, newkp, label=label, **kwargs)
applySmartTimeTicks(ax, time_range, dolabel=True)
return fig, ax
def create_figure(iefile):
# Open IE file, calculate azimuthal current:
ie = rim.Iono(iefile)
ie.calc_j()
ie['n_jphi'] /= 1000.
ie['s_jphi'] /= 1000.
# Create figure:
fig = plt.figure(figsize=(10, 10))
fig.subplots_adjust(left=.08,
bottom=.052,
top=.95,
right=.962,
wspace=.217,
hspace=.23)
# Set common keyword arguments:
kwargs = {'max_colat': args.colat, 'target': fig, 'add_cbar': True}
# Add IE plots:
# FACs:
with plt.style.context('default'):
zmax = max(np.abs(ie['n_jr']).max(), np.abs(ie['s_jr']).max())
ie.add_cont('n_jr', loc=431, **kwargs)
ie.add_cont('s_jr', loc=434, label='', **kwargs)
# J_phi:
zmax = max(np.abs(ie['n_jphi']).max(), np.abs(ie['s_jphi']).max())
ie.add_cont('n_jphi', loc=432, **kwargs)
ie.add_cont('s_jphi', loc=435, label='', **kwargs)
# J_phi:
zmax = max(np.abs(ie['n_phi']).max(), np.abs(ie['s_phi']).max())
ie.add_cont('n_phi', loc=433, **kwargs)
ie.add_cont('s_phi', loc=436, label='', **kwargs)
a4 = fig.add_subplot(413)
a5 = fig.add_subplot(414, sharex=a4)
# Add IMF data:
a4.plot(imf['time'], imf['bz'], color=bzcolor, lw=1.25)
a4.plot(imf['time'], imf['by'], color=bycolor, lw=1.25)
a4.hlines(0,
imf['time'][0],
imf['time'][-1],
color='k',
linestyles='dashed')
a4.legend(['B$_Z$', 'B$_Y$'], loc='upper left')
a4.set_ylabel('IMF ($nT$)')
# Add Pdyn:
a5.plot(imf['time'], imf['pram'], color=pdcolor)
a5.grid(False, axis='y')
a5.set_ylabel('P$_{dyn}$ ($nPa$)')
applySmartTimeTicks(a5, imf['time'], dolabel=True)
a4.vlines(ie.attrs['time'],
a4.get_ylim()[0],
a4.get_ylim()[1],
colors='k',
linestyles='dashed',
linewidths=2)
lim = a5.get_ylim()
a5.vlines(ie.attrs['time'],
lim[0],
lim[1],
colors='k',
linestyles='dashed',
linewidths=2)
fig.savefig(outdir + f'iono_t{ie.attrs["time"]:%Y%m%d_%H%M%S}.png')
return fig
def plot_indexes():
# Some constants:
bzcolor = '#3333CC'
bycolor = '#ff9900'
pdcolor = '#CC3300'
# Find log files:
log_path = glob('./GM/log*.log')
ind_path = glob('./GM/geoind*.log')
# Get that IMF file:
imfpath = glob('./imf*.dat')[0] if not args.imffile else args.imffile
imf = ImfInput(imfpath)
if not (log_path) or not (ind_path):
raise ValueError('Cannot find log or geoindex file in GM/')
# Open index/log files:
log = bats.BatsLog(log_path[0])
ind = bats.GeoIndexFile(ind_path[0])
# Create figure and axes:
f1 = plt.figure(figsize=(8.5, 11))
a1, a2, a3, a4, a5 = f1.subplots(5, 1, sharex=True)
# Set line style for observations:
obs_kwargs = {'c': 'k', 'ls': '--', 'alpha': .7}
# Dst:
plot1 = log.add_dst_quicklook(target=a1,
plot_sym=args.obs,
obs_kwargs=obs_kwargs)
# Kp:
plot2 = ind.add_kp_quicklook(target=a2, loc=412, plot_obs=args.obs)
# Ae:
plot3 = ind.add_ae_quicklook(target=a3,
plot_obs=args.obs,
val='AU',
obs_kwargs=obs_kwargs)
plot3 = ind.add_ae_quicklook(target=a3,
plot_obs=args.obs,
val='AL',
obs_kwargs=obs_kwargs,
c=a3.get_lines()[0].get_color())
# Replace legend with something better.
a3.set_ylabel('AU/AL ($nT$)')
a3.legend(a3.lines[:2], ['Model AU/AL', 'Obs. AU/AL'], loc='best')
# Add IMF data:
a4.plot(imf['time'], imf['bz'], color=bzcolor, lw=1.25)
a4.plot(imf['time'], imf['by'], color=bycolor, lw=1.25)
a4.hlines(0,
ind['time'][0],
ind['time'][-1],
color='k',
linestyles='dashed')
a4.legend(['B$_Z$', 'B$_Y$'], loc='best')
a4.set_ylabel('IMF ($nT$)')
# Add Pdyn:
a5.plot(imf['time'], imf['pram'], color=pdcolor)
a5.grid(False, axis='y')
a5.set_ylabel('P$_{dyn}$ ($nPa$)')
# Set axes ticks and x-labels:
applySmartTimeTicks(a5, ind['time'], dolabel=True)
for a in [a1, a2, a3, a4]:
a.set_xlabel('')
#a.set_xticklabels('')
a1.set_title('Geomagnetic Indices: {:%Y-%m-%d}'.format(log['time'][0]),
size=20)
plt.tight_layout()
f1.savefig('./geoindexes.png')
def plot_iono():
'''
Create a plot to summarize ionospheric activity.
Requires total current calculation!
'''
from pickle import load
# Some constants:
bzcolor = '#3333CC'
bycolor = '#ff9900'
pdcolor = '#CC3300'
north, south = '#0064b1', '#f58025'
# Find log files:
log_path = glob('./IE/IE*.log')
ind_path = glob('./GM/geoind*.log')
if not (log_path) or not (ind_path):
raise ValueError('Cannot find log or geoindex file in IE/ and GM/')
# Open currents:
# time, nUp, nDown, nAll, nPhiMax, nPhiMin, sUp, sDown, sAll, sPhiMax, sPhiMin
current_file = glob('./*.pkl')[0] if not args.currents else args.currents
with open(current_file, 'rb') as f:
curr = load(f)
# Get that IMF file:
imfpath = glob('./imf*.dat')[0] if not args.imffile else args.imffile
imf = ImfInput(imfpath)
# Open index/log files:
log = bats.BatsLog(log_path[0])
ind = bats.GeoIndexFile(ind_path[0])
# Create figure and axes:
f1 = plt.figure(figsize=(8.5, 11))
a1, a2, a3, a4, a5 = f1.subplots(5, 1, sharex=True)
# CPCP:
a1.plot(log['time'], log['cpcpn'], c=north, lw=2., label='Northern Hemi.')
a1.plot(log['time'], log['cpcps'], c=south, lw=2., label='Southern Hemi.')
a1.set_ylabel('CPCP ($kV$)')
a1.legend(loc='best')
# Total current:
a2.plot(curr[0], curr[3], c=north, lw=2.)
a2.plot(curr[0], curr[8], c=south, lw=2.)
a2.set_ylabel('Total J$_{\parallel}$ ($MA$)')
# Peak azimuthal:
a3.plot(curr[0], curr[4] / 1000., c=north, lw=2.)
a3.plot(curr[0], curr[5] / 1000., c=north, lw=2.)
a3.plot(curr[0], curr[9] / 1000., c=south, lw=2.)
a3.plot(curr[0], curr[10] / 1000., c=south, lw=2.)
a3.set_ylabel('Max/Min J$_{\\Phi}$ ($mA/m$)')
# Add IMF data:
a4.plot(imf['time'], imf['bz'], color=bzcolor, lw=2)
a4.plot(imf['time'], imf['by'], color=bycolor, lw=2)
a4.hlines(0,
ind['time'][0],
ind['time'][-1],
color='k',
linestyles='dashed')
a4.legend(['B$_Z$', 'B$_Y$'], loc='best')
a4.set_ylabel('IMF ($nT$)')
# Add Pdyn:
a5.plot(imf['time'], imf['pram'], color=pdcolor)
a5.grid(False, axis='y')
a5.set_ylabel('P$_{dyn}$ ($nPa$)')
# Set axes ticks and x-labels:
applySmartTimeTicks(a5, ind['time'], dolabel=True)
for a in [a1, a2, a3, a4]:
a.set_xlabel('')
#a.set_xticklabels('')
a1.set_title('Ionospheric Summary: {:%Y-%m-%d}'.format(log['time'][0]),
size=20)
plt.tight_layout()
f1.savefig('./iono_summary.png')
def add_histplot(self,
target=False,
loc=111,
label='Kyoto $K_{p}$',
time_range=None,
**kwargs):
'''
Make a quick histogram-style plot of the Kp data.
Returns
=======
fig : matplotlib figure object
ax : matplotlib axes object
Other Parameters
================
target : Figure or Axes
If None (default), a new figure is generated from scratch.
If a matplotlib Figure object, a new axis is created
to fill that figure.
If a matplotlib Axes object, the plot is placed
into that axis.
loc : int
Use to specify the subplot placement of the axis
(e.g. loc=212, etc.) Used if target is a Figure or None.
Default 111 (single plot).
label : string
The label applied to the line when a legend is added to the axes.
Defaults to 'Kyoto $K_{p}$'.
time_range : tuple of datetimes
The time range to plot. Only the first and last values in the
tuple (or list) are used if the number of elements is greater than
two. Defaults to **None**, meaning that the full time range
available is used.
Extra keyword arguments are passed to :function:`matplotlib.pyplot.plot`
to customize the line style.
Examples
========
>>> import spacepy.pybats.kyoto as kt
>>> kp = kt.fetch('kp', (1981, 11), (1981, 11)
>>> kp.add_histplot(lw=2.0, lc='r', label='Example KP')
'''
import matplotlib.pyplot as plt
# Shortcuts for the lazy.
bstart = self['binstart']
bstop = self['binstop']
npts = self.attrs['npts']
# Reformulate time to get histogram-type look.
newtime = np.zeros(npts * 24, dtype=object)
newkp = np.zeros(npts * 24)
# Set time range.
if not time_range:
time_range = newtime
for i in range(npts * 8):
newtime[3 * i] = bstart[i]
newtime[3 * i + 1] = self['time'][i]
newtime[3 * i + 2] = bstop[i]
newkp[3 * i:3 * i +
3] = self['kp'][i], self['kp'][i], self['kp'][i]
fig, ax = set_target(target, figsize=(10, 4), loc=loc)
line = ax.plot(newtime, newkp, label=label, **kwargs)
applySmartTimeTicks(ax, time_range, dolabel=True)
return fig, ax
def plotSummary(self, timerange=None, coord_sys=None, fig_target=None, spec=False, orbit_params=(False, True),
**kwargs):
"""Generate summary plot of AE9/AP9/SPM data loaded
spec : if True, plot spectrogram instead of flux/fluence lineplot, requires 'ecol' keyword
"""
if timerange:
if isinstance(timerange, spt.Ticktock):
t1 = timerange.UTC[0]
t2 = timerange.UTC[-1]
elif isinstance(timerange[0], dt.datetime):
t1 = timerange[0]
t2 = timerange[-1]
else:
raise TypeError('Incorrect data type provided for timerange')
# now select subset
i_use = tb.tOverlapHalf([t1, t2], self['Epoch'])
t_use = self['Epoch'][i_use]
c_use = self['Coords'][i_use]
f_use = self[self.attrs['varname']][i_use, ...]
else:
t_use = self['Epoch']
c_use = self['Coords']
f_use = self[self.attrs['varname']]
if coord_sys and (coord_sys.upper() != c_use.attrs['COORD_SYS']):
# TODO: We assume cartesian, make flexible so can take spherical
cx = spc.Coords(c_use, c_use.attrs['COORD_SYS'], 'car')
cx.ticks = spt.Ticktock(t_use)
cx = cx.convert(coord_sys.upper(), 'car').data
else:
coord_sys = c_use.attrs['COORD_SYS']
cx = c_use
sys_subs = r'$_{' + coord_sys + r'}$'
splot.style('spacepy')
if orbit_params[0]:
landscape = orbit_params[1]
locs = [121, 122] if landscape else [211, 212]
else:
locs = [223, 224]
landscape = True
if fig_target:
fig, ax1 = splot.set_target(fig_target, loc=locs[0])
else:
fig, ax1 = splot.set_target(fig_target, figsize=(8, 8), loc=locs[0])
fig, ax2 = splot.set_target(fig, loc=locs[1])
ax1 = self._makeOrbitAxis(cx[:, 0], cx[:, 1], ax1)
ax2 = self._makeOrbitAxis(cx[:, 0], cx[:, 2], ax2)
if landscape: ax1.set_xlabel('X{0} [{1}]'.format(sys_subs, self['Coords'].attrs['UNITS']))
ax2.set_xlabel('X{0} [{1}]'.format(sys_subs, self['Coords'].attrs['UNITS']))
ax1.set_ylabel('Y{0} [{1}]'.format(sys_subs, self['Coords'].attrs['UNITS']))
ax2.set_ylabel('Z{0} [{1}]'.format(sys_subs, self['Coords'].attrs['UNITS']))
ax1xl, ax1yl, ax2xl, ax2yl = ax1.get_xlim(), ax1.get_ylim(), ax2.get_xlim(), ax2.get_ylim()
maxabslim = np.max(np.abs([ax1xl, ax1yl, ax2xl, ax2yl]))
if np.abs((ax1.get_xlim()[1] - ax1.get_xlim()[0])) < 1:
refpt = maxabslim
ax1.set_xlim([-1.25 * refpt, 1.25 * refpt])
ax1.set_ylim(ax1.get_xlim())
refpt = ax2.get_xlim()[0]
ax2.set_xlim([-1.25 * refpt, 1.25 * refpt])
ax2.set_ylim(ax1.get_xlim())
else:
ax1.set_xlim([-maxabslim, maxabslim])
ax1.set_ylim([-maxabslim, maxabslim])
ax2.set_xlim([-maxabslim, maxabslim])
ax2.set_ylim(ax2.get_xlim())
ax1.invert_yaxis()
ax1.invert_xaxis()
ax2.invert_xaxis()
ax1.set_aspect('equal')
ax2.set_aspect('equal')
if not orbit_params[0]:
ax3 = splot.plt.subplot2grid((2, 2), (0, 0), colspan=2)
if not spec:
l3 = ax3.semilogy(t_use, f_use)
ylab = '{0} ['.format(self.attrs['varname']) + re.sub('(\^[\d|-]*)+', _grp2mathmode,
self[self.attrs['varname']].attrs['UNITS']) + ']'
ax3.set_ylabel(ylab)
for ll, nn in zip(l3, self['Energy']):
ll.set_label('{0} {1}'.format(nn, self['Energy'].attrs['UNITS']))
ncol = len(self['Energy']) // 2 if len(self['Energy']) <= 6 else 3
leg = ax3.legend(loc='center', bbox_to_anchor=(0.5, 1), ncol=ncol,
frameon=True, framealpha=0.5)
lims3 = ax3.get_ylim()
newupper = 10 ** (np.log10(lims3[0]) + (np.log10(lims3[1] / lims3[0]) * 1.125))
ax3.set_ylim([lims3[0], newupper])
splot.applySmartTimeTicks(ax3, t_use)
fig.tight_layout()
pos3 = ax3.get_position()
ax3.set_position([pos3.x0, pos3.y0, pos3.width, pos3.height * 0.8])
# fig.suptitle('{model_type}\n'.format(**self.attrs) +
# '{0} - {1}'.format(t_use[0].isoformat()[:19], t_use[-1].isoformat()[:19]))
else:
if timerange: raise NotImplementedError('Time range selection not yet implemented for spectrograms')
pos3 = ax3.get_position()
ax3.set_position([pos3.x0, pos3.y0, pos3.width, pos3.height * 0.9])
ecol = kwargs['ecol'] if 'ecol' in kwargs else 0
ax3 = self.plotSpectrogram(target=ax3, ecol=ecol)
splot.plt.subplots_adjust(wspace=0.3)
pos1 = ax1.get_position()
ax1.set_position([pos3.x0, pos1.y0, pos1.width, pos1.height])
splot.revert_style()
return fig
def plot_results(fileY, fileZ, iFile=0, nFiles=1):
'''
Create a nice plot of WTF is happening.
Kwargs iFile and nFiles are for reporting progress.
'''
# Get start times from file names:
time_y = pb.parse_filename_time(fileY)
time_z = pb.parse_filename_time(fileZ)
# Check that times match:
for t1, t2 in zip(time_y, time_z):
if t1 != t2:
print(
'File times do not match!\n\tY-File: {}\n\tZ-File: {}'.format(
t1, t2))
raise ValueError('File time mismatch')
# Use file time info get actual time:
if time_y[-1]:
tNow = time_y[-1] # Datetime in file name.
elif time_y[1]:
t_delta = dt.timedelta(seconds=time_y[1])
tNow = start + t_delta # Runtime in file name.
else:
tNow = start # Only iteration given.
# Check if files are outside time range.
if tNow < trange[0]: return 0
if tNow > trange[-1]: return 0
# Open files & calculate important values:
mY = pbs.Bats2d(fileY)
mZ = pbs.Bats2d(fileZ)
mY.calc_ndens()
mZ.calc_ndens()
mY.calc_utotal()
mZ.calc_utotal()
if mZ.attrs['runtime'] != mY.attrs['runtime']:
print('{}\n{}\n'.format(fileY, fileZ))
raise (ValueError('Times of files do not line up!'))
fig = plt.figure(figsize=[11, 8.5])
fig.subplots_adjust(left=0.08, right=.96, bottom=0.06, top=.93)
a1 = fig.add_subplot(221)
stuff = mY.add_contour('x',
'z',
v1,
91,
dolog=True,
xlim=xrng,
ylim=zrng,
add_cbar=True,
target=a1,
zlim=prng,
cmap=cmap1)
a1.set_title('{} ({})'.format(pb.mhdname_to_tex(v1),
mY[v1].attrs['units']))
stuff[-1].set_label('')
if args.mag: mY.add_b_magsphere(a1, DoOpen=True)
a2 = fig.add_subplot(222)
stuff = mZ.add_contour('x',
'y',
v2,
91,
dolog=True,
xlim=xrng,
ylim=yrng,
add_cbar=True,
target=a2,
zlim=drng,
cmap=cmap2)
a2.set_title('{} ({})'.format(pb.mhdname_to_tex(v2),
mZ[v2].attrs['units']))
stuff[-1].set_label('')
for s in sats:
s.add_sat_loc(tNow, a1, plane='XZ', dolabel=True, color='r', size=9)
s.add_sat_loc(tNow, a2, plane='XY', dolabel=True, color='r', size=9)
pos = stuff[-1].ax.get_position()
a3, a4 = fig.add_subplot(413), fig.add_subplot(414)
a5 = a4.twinx()
pos3, pos4 = a3.get_position(), a4.get_position()
a3.set_position([pos3.xmin, pos3.ymin, pos3.width * .95, pos3.height])
a4.set_position([pos4.xmin, pos4.ymin, pos4.width * .95, pos4.height])
a5.set_position([pos4.xmin, pos4.ymin, pos4.width * .95, pos4.height])
log.add_dst_quicklook(target=a3,
plot_obs=args.obsdst,
obs_kwargs={
'c': 'k',
'ls': '-'
},
color='r')
a3.legend(loc='best', fontsize=10)
a3.set_xlabel('')
applySmartTimeTicks(a3, trange, dolabel=False)
a3.grid(False, axis='y')
a3.set_ylabel('D$_{ST}$ ($nT$)')
a3.hlines(0, trange[0], trange[1], color='grey', linestyles='dashed')
ymin, ymax = a3.get_ylim()
a3.vlines(tNow, ymin, ymax, linestyles='solid', color='g', linewidths=2.)
a3.set_ylim([ymin, ymax])
a4.plot(imf['time'], imf['bz'], color=bzcolor, label='B$_Z$')
if args.by:
a4.plot(imf['time'], imf['by'], color=bycolor, label='B$_Y$')
a4.legend(loc='best', ncol=2)
a5.plot(imf['time'], imf['pram'], color=pdcolor, lw=1.5)
applySmartTimeTicks(a4, trange, dolabel=True)
a4.grid(False, axis='y')
a4.hlines(0, trange[0], trange[1], color=bzcolor, linestyles='dashed')
a4.set_ylabel('IMF ' + (not args.by) * 'B$_{Z}$ ' + '($nT$)',
color=bzcolor)
a5.set_ylabel('P$_{dyn}$ ($nPa$)', color=pdcolor)
a5.grid(False, axis='y')
ymin, ymax = a4.get_ylim()
a4.vlines(tNow, ymin, ymax, linestyles='solid', color='g', linewidths=2.)
a4.set_ylim([ymin, ymax])
fig.suptitle(args.title, size=18)
# Save the figure (normal operation) or show figure (debug mode).
if args.debug:
plt.show()
else:
fig.savefig(args.outdir + '/mhd{:%Y%m%d_%H%M%S}.png'.format(tNow))
# Clear and close figure.
if not args.debug:
fig.clf()
plt.close(fig)
percent_done = float(iFile + 1) / float(nFiles) * 100.0
return percent_done
