def plotMesh(xx, yy, zz, cmap='plasma', target=None, horiz=False):
fig, ax = splot.set_target(target)
orient = 'horizontal' if horiz else 'vertical'
map_ax = ax.pcolormesh(xx, yy, zz, cmap=cmap)
cax = fig.colorbar(map_ax, ax=ax, orientation=orient)
m0 = 0 # colorbar min value
ntens = 10 if (1 / cax.vmax <= 10) else 100
m5 = roundUpTo(cax.vmax, ntens) # colorbar max value
num_ticks = 5
newticks = np.linspace(m0, m5, num_ticks)
cax.set_ticks(newticks.tolist())
cax.set_ticklabels(['{:.2f}'.format(x) for x in newticks.tolist()])
return ax, cax
def plotlog(self, vars, xlim=None, ylim=None,
target=None,
loc=111, title=None,
dolog=False,
figsize=(8.34,7),
*args, **kwargs):
'''
Create a plot of variable *var* at a given time.
If kwarg **target** is None (default), a new figure is
generated from scratch. If target is a matplotlib Figure
object, a new axis is created to fill that figure at subplot
location **loc**. If **target** is a matplotlib Axes object,
the plot is placed into that axis.
'''
import matplotlib.pyplot as plt
fig, ax = set_target(target, loc=loc, figsize=figsize)
#ax.set_aspect('equal')
# Grab values from correct time/location.
for var in vars:
x = self['time'][:]
y = self[var][:]
ax.plot(x,y,label=var)
# Create levels and set norm based on dolog.
if dolog:
plt.yscale('log')
#add legend
ax.legend()
if xlim != None:
ax.set_xlim(xlim[0],xlim[1])
if ylim != None:
ax.set_ylim(ylim[0],ylim[1])
return fig, ax
def add_contour(self,
var,
zlim=None,
target=None,
loc=111,
title=None,
Lmax=None,
add_cbar=False,
clabel=None,
dolog=False,
filled=True,
nLev=31,
labelsize=14,
**kwargs):
'''
Create a polar contour plot of variable *var* and add it to *target*.
Extra keyword arguments are handed to matplotlib's contourf command.
Parameters
==========
var : string
The variable within the object to plot.
Returns
=======
fig : matplotlib figure object
ax : matplotlib axes object
cont : matplotlib contour object
cbar : matplotlib colorbar object
Other Parameters
================
nLev : int
Sets the number of contour levels. Defaults to 31.
zlim : two-element list or array
Set the color bar range. Some variables have default ranges. Others
will use max and min of *var*.
Lmax : real
Set the radial extent of the plot.
add_cbar : bool
Set whether to add a color bar or not. Defaults to **False**.
dolog : bool
If **True**, use a log scale to plot *var*. Defaults to **False**.
labelsize : int
Sets the font size of the labels. Defaults to 14.
title : string
Sets the plot title. Defaults to 'auto', using the variable label.
clabel : string
Set label for the color bar. Defaults to *var* and associated units.
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).
'''
from numpy import linspace, pi
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm
from matplotlib.ticker import MultipleLocator, LogLocator
# Set ax and fig based on given target.
fig, ax = set_target(target, figsize=(10.5, 8), loc=loc, polar=True)
# Set function based on boolean "filled":
if filled:
contour = ax.contourf
else:
contour = ax.contour
# Get max/min if none given.
if zlim == None:
if var in self._default_zlims:
zlim = self._default_zlims[var]
else:
zlim = [self[var].min(), self[var].max()]
if dolog and zlim[0] <= 0:
zlim[0] = np.min([0.0001, zlim[1] / 1000.0])
# Create levels and set norm based on dolog.
if dolog:
levs = np.power(
10, np.linspace(np.log10(zlim[0]), np.log10(zlim[1]), nLev))
z = np.where(self[var] > zlim[0], self[var], 1.01 * zlim[0])
z[z > zlim[-1]] = zlim[-1]
norm = LogNorm()
lct = LogLocator()
else:
levs = np.linspace(zlim[0], zlim[1], nLev)
z = self[var]
norm = None
lct = None
# Allow results to cross phi=360.
phi = np.concatenate(
(self['lon'], [360.0])) * np.pi / 180. - np.pi / 2.
z = np.concatenate((z, np.array([z[:, 0]]).transpose()), 1)
# Set default color tables based on variable plotted.
if ('cmap' not in kwargs) and ('colors' not in kwargs) and \
(var in self._default_cmaps):
kwargs['cmap'] = self._default_cmaps[var]
# -------Plot away------
# Contour:
cont = contour(phi, self['L'], z, levs, norm=norm, **kwargs)
# Add cbar if requested:
if add_cbar:
cbar = plt.colorbar(cont, pad=0.08, shrink=.8, ticks=lct, ax=ax)
if clabel == None:
clabel = "%s ($%s$)" % (var, self[var].attrs['units'])
cbar.set_label(clabel)
else:
cbar = None
# Adjust plot appropriately.
if not Lmax:
# Default to inside ghost cells.
Lmax = self['L'][-3]
if title:
ax.set_title(title + '\n' + self.attrs['time'].isoformat(),
position=(0, 1),
ha='left',
size=14)
_adjust_dialplot(ax, Lmax, labelsize=labelsize)
return fig, ax, cont, cbar
def add_pcolor(self,
var,
zlim=None,
target=None,
loc=111,
title=None,
Lmax=None,
add_cbar=False,
clabel=None,
dolog=False,
labelsize=14,
**kwargs):
'''
Create a polar pcolor plot of variable *var* and add it to *target*.
Extra keyword arguments are handed to matplotlib's pcolor command.
Parameters
==========
var : string
The variable within the object to plot.
Returns
=======
fig : matplotlib figure object
ax : matplotlib axes object
cont : matplotlib contour object
cbar : matplotlib colorbar object
Other Parameters
================
zlim : two-element list or array
Set the color bar range. Some variables have default ranges. Others
will use max and min of *var*.
Lmax : real
Set the radial extent of the plot.
add_cbar : bool
Set whether to add a color bar or not. Defaults to **False**.
dolog : bool
If **True**, use a log scale to plot *var*. Defaults to **False**.
labelsize : int
Sets the font size of the labels. Defaults to 14.
title : string
Sets the plot title. Defaults to 'auto', using the variable label.
clabel : string
Set label for the color bar. Defaults to *var* and associated units.
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).
'''
from numpy import linspace, pi
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm
# Set ax and fig based on given target.
fig, ax = set_target(target, figsize=(10.5, 8), loc=loc, polar=True)
# Get max/min if none given.
if zlim == None:
if var in self._default_zlims:
zlim = self._default_zlims[var]
else:
zlim = [self[var].min(), self[var].max()]
if dolog and zlim[0] <= 0:
zlim[0] = np.min([0.0001, zlim[1] / 1000.0])
# Logarithmic scale?
if dolog:
z = np.where(self[var] > zlim[0], self[var], 1.01 * zlim[0])
norm = LogNorm()
else:
z = self[var]
norm = None
# Set up proper meshes.
nL, nPhi = len(self['L']), len(self['lon'])
dL, dPhi = self.attrs['dL'], self.attrs['dLon'] * pi / 180.0
phi = linspace(-1.0 * dPhi / 2.0, 2. * pi - dPhi / 2.0,
nPhi + 1) - pi / 2.
r = linspace(self['L'][0] - dL / 2.0, self['L'][-1] + dL / 2.0, nL + 1)
# Set default color tables based on variable plotted.
if ('cmap' not in kwargs) and var in self._default_cmaps:
kwargs['cmap'] = self._default_cmaps[var]
# -------Plot away------
# Mesh:
pcol = ax.pcolormesh(phi,
r,
z,
vmin=zlim[0],
vmax=zlim[1],
norm=norm,
**kwargs)
# Add cbar if requested:
if add_cbar:
cbar = plt.colorbar(pcol, pad=0.08, shrink=.8, ax=ax)
if clabel == None:
clabel = "%s ($%s$)" % (var, self[var].attrs['units'])
cbar.set_label(clabel)
else:
cbar = None
# Adjust plot appropriately.
if not Lmax:
# Default to inside ghost cells.
Lmax = self['L'][-3]
if title:
ax.set_title(title + '\n' + self.attrs['time'].isoformat(),
position=(0, 1),
ha='left',
size=14)
_adjust_dialplot(ax, Lmax, labelsize=labelsize)
return fig, ax, pcol, cbar
def add_oval_line(self, time, *args, **kwargs):
'''
Adds the location of the auroral oval at time *time* as a line plot on to
*target*, which must be a Matplotlib figure/axes.
If *target* not given, a new axes object is created. If *target* is
not an existing axes object, a new axes object is created and customized
to be an ionosphere polar plot.
Extra arguments/kwargs are sent to :meth:`matplotlib.pyplot.plot`.
Parameters
==========
time : integer or datetime
Sets the time at which to plot the line. If an integer is used,
the integer is used to index the internal array naively. If a
datetime object is given and is within the bounds of self['time'],
the oval location will be interpolated to *time*. Control of the
interpolation happens via the *interp* keyword.
Returns
=======
fig : matplotlib figure object
ax : matplotlib axes object
line : matplotlib line 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).
interp : bool
Control the behavior of time interpolation when the *time* argument
is given as a datetime object. Defaults to **True**, meaning that
oval location is interpolated in time.
Examples
========
>>>
'''
import datetime as dt
# Handle kwargs not being handed to plot:
target=None; loc=111; interp=True
if 'target' in kwargs: target=kwargs.pop('target')
if 'interp' in kwargs: interp=kwargs.pop('interp')
if 'loc' in kwargs: loc=kwargs.pop('loc')
# Set plot targets:
fig, ax = set_target(target, polar=True, loc=loc)
# Get oval interpolated in time:
if type(time) == dt.datetime:
oval = self.get_oval(time, interp=interp)
elif type(time) == int:
oval = self['oval'][time]
else:
raise TypeError('Unrecognized type for *time*:'+type(time))
# Plot oval:
line = ax.plot(self['lon']+np.pi/2., oval, *args, **kwargs)
return fig, ax, line
outputlist = inputlist/fac
except TypeError: #not an array
outputlist = [el/fac for el in inputlist]
return outputlist
isi = convert_tdelt(results['tau_valid'], units='minutes')
isi_hr = convert_tdelt(results['tau_valid'], units='hours')
tau_ax = np.arange(0,30*60,.2)
try:
from sklearn.neighbors.kde import KernelDensity
kern_type = 'epanechnikov'
kern_lab = '{0}{1} KDE'.format(kern_type[0].upper(), kern_type[1:])
kernel = KernelDensity(kernel=kern_type, bandwidth=60).fit(isi[:, np.newaxis])
kde_plot = np.exp(kernel.score_samples(tau_ax[:, np.newaxis]))
except ImportError:
from scipy import stats
kern_lab = 'Gaussian KDE'
kernel = stats.gaussian_kde(isi, bw_method='scott')
kde_plot = kernel.evaluate(tau_ax)
fig, ax = splot.set_target(None)
ax.hist(isi_hr, bins=np.arange(0,25,0.5), histtype='step', normed=True, lw=1.5, label='Binned Data')
ax.plot(tau_ax/60., kde_plot*60., lw=1.5, label=kern_lab)
ax.set_xlim([0,25])
ax.set_ylabel('Probability')
ax.set_xlabel(r'Inter-substorm Interval, $\tau$ [hours]') #raw string req'd (else \t in \tau becomes [tab]au
ax.legend()
fig.suptitle('MSM$_{Python}$: ' + '{0} (1998-2002)'.format(satname))
plt.show()
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 add_slice(self,
var,
time,
nlev=31,
zlim=None,
target=None,
loc=111,
title=None,
latoffset=1.05,
rlim=50.,
add_cbar=True,
clabel=None,
show_pts=False,
dolog=False,
lats=[75., 60.],
colats=None,
figsize=(8.34, 7),
*args,
**kwargs):
'''
Create a plot of variable *var*.
If kwarg **target** is None (default), a new figure is
generated from scratch. If target is a matplotlib Figure
object, a new axis is created to fill that figure at subplot
location **loc**. If **target** is a matplotlib Axes object,
the plot is placed into that axis.
'''
import matplotlib.pyplot as plt
from matplotlib.colors import (LogNorm, Normalize)
from matplotlib.patches import Circle
from matplotlib.ticker import (LogLocator, LogFormatter,
LogFormatterMathtext, MultipleLocator)
# Grab the slice of data that we want:
if type(time) == type(self['time'][0]):
if time not in self['time']:
raise ValueError('Time not in object')
time = np.arange(self.attrs['nTime'])[self['time'] == time][0]
fig, ax = set_target(target, loc=loc, figsize=figsize)
ax.set_aspect('equal')
# Create values in Cartesian plane.
self._get_latlon()
# Grab values from correct time/location.
x = self._xLat[time, :]
y = self._yLat[time, :]
value = self[var][time, :]
# Get max/min if none given.
if zlim == None:
zlim = [0, 0]
zlim[0] = value.min()
zlim[1] = value.max()
if dolog and zlim[0] <= 0:
zlim[0] = np.min([0.0001, zlim[1] / 1000.0])
# Create levels and set norm based on dolog.
if dolog:
levs = np.power(
10, np.linspace(np.log10(zlim[0]), np.log10(zlim[1]), nlev))
z = np.where(value > zlim[0], value, 1.01 * zlim[0])
norm = LogNorm()
ticks = LogLocator()
fmt = LogFormatterMathtext()
else:
levs = np.linspace(zlim[0], zlim[1], nlev)
z = value
norm = None
ticks = None
fmt = None
# Create contour plot.
cont=ax.tricontourf(np.asarray(x), np.asarray(y), np.asarray(z), \
np.asarray(levs), *args, norm=norm, **kwargs)
if show_pts:
ax.plot(x, y, '+w')
# Add cbar if necessary.
if add_cbar:
cbar = plt.colorbar(cont, ticks=ticks, format=fmt, pad=0.01)
if clabel == None:
clabel = "%s (%s)" % (var, self[var].attrs['units'])
cbar.set_label(clabel)
else:
cbar = None # Need to return something, even if none.
# Set title, labels, axis ranges (use defaults where applicable.)
if title: ax.set_title(title)
ax.set_yticks([]), ax.set_xticks([])
ax.set_ylabel(r'Sun $\rightarrow$')
colat_lim = 90. - rlim
ax.set_xlim([-1 * colat_lim, colat_lim])
ax.set_ylim([-1 * colat_lim, colat_lim])
if colats:
for colat in colats:
r = latoffset * colat / np.sqrt(2)
ax.add_patch(
Circle([0, 0], colat, fc='none', ec='k', ls='dashed'))
ax.text(r,
r,
'{:.0f}'.format(colat) + r'$^{\circ}$',
rotation=315,
ha='center',
va='center')
else:
for lat in lats:
colat = 90 - lat
r = latoffset * colat / np.sqrt(2)
ax.add_patch(
Circle([0, 0], colat, fc='none', ec='k', ls='dashed'))
ax.text(r,
r,
'{:.0f}'.format(lat) + r'$^{\circ}$',
rotation=315,
ha='center',
va='center')
return fig, ax, cont, cbar
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_separatrix(self, target=None,loc=111,Lmax=None,title=None,**kwargs):
'''
Attempts to locate the separatrix (separator between closed and
open drift paths) by finding the minimum velocity in the domain and
tracing the path of constant potential that passes through that point.
The figure, axes, and contour object containing the separatrix line
are returned to the user.
If kwarg **target** is None (default), a new figure is
generated from scratch. If target is a matplotlib Figure
object, a new axis is created to fill that figure at subplot
location **loc**. If **target** is a matplotlib Axes object,
the plot is placed into that axis.
Four values are returned: the matplotlib Figure and Axes objects,
the matplotlib contour object, and the matplotlib colorbar object
(defaults to *False* if not used.)
Kwargs that set line characteristics behave in the typical matplotlib
manner (i.e. "colors" can be set to either a color name or a
hexidecimal specifier.)
=========== ==========================================================
Kwarg Description
----------- ----------------------------------------------------------
target Set plot destination. Defaults to new figure.
loc Set subplot location. Defaults to 111.
linewidths Set width of plotted line. Defaults to 3.0.
colors Set color of line. Defaults to 'orange'.
linestyles Set line style. Defaults to 'dashed'.
----------- ----------------------------------------------------------
'''
from numpy import linspace, pi, sqrt
import matplotlib.pyplot as plt
# Set default line behavior.
if 'linewidths' not in kwargs:
kwargs['linewidths'] = 3.0
if 'colors' not in kwargs:
kwargs['colors'] = 'orange'
if 'linestyles' not in kwargs:
kwargs['linestyles'] = 'solid'
# Set ax and fig based on given target.
fig, ax = set_target(target, figsize=(10.5,8), loc=loc, polar=True)
doAdjust = not target==ax
# Find the stagnation point by looking for E = 0.
# Get the potential value at that spot.
# Exclude near-boundary points and limit azimuthal extent.
if 'E' not in self:
self.calc_E()
i = np.arange(self['L'].size )[self['L']<9.75]
j = np.arange(self['lon'].size)[(self['lon']>250.) & (self['lon']<360.)]
pmin = self['pot'][self['E'] == self['E'][i[0]:i[-1],j[0]:j[-1]].min()][0]
# Create a contour that only follows that potential curve.
phi = np.concatenate( (self['lon'], [360.0]) ) * pi/180. - pi/2.
z = np.concatenate( (self['pot'],
np.array([self['pot'][:,0]]).transpose() ), 1)
cnt = ax.contour(phi, self['L'], z, levels=[pmin], **kwargs)
# Adjust plot appropriately.
if doAdjust:
if not Lmax:
# Default to inside ghost cells.
Lmax = self['L'][-3]
if title:
ax.set_title(title+'\n'+self.attrs['time'].isoformat(),
position=(0,1), ha='left', size=14)
_adjust_dialplot(ax, Lmax, labelsize=14)
# Return bits to user.
return fig, ax, cnt
def add_contour(self, var, zlim=None, target=None, loc=111, title=None,
Lmax=None, add_cbar=False, clabel=None, dolog=False,
filled=True, nLev=31, labelsize=14, **kwargs):
'''
Create a polar contour plot of variable *var* and add it to *target*.
Extra keyword arguments are handed to matplotlib's contourf command.
Parameters
==========
var : string
The variable within the object to plot.
Returns
=======
fig : matplotlib figure object
ax : matplotlib axes object
cont : matplotlib contour object
cbar : matplotlib colorbar object
Other Parameters
================
nLev : int
Sets the number of contour levels. Defaults to 31.
zlim : two-element list or array
Set the color bar range. Some variables have default ranges. Others
will use max and min of *var*.
Lmax : real
Set the radial extent of the plot.
add_cbar : bool
Set whether to add a color bar or not. Defaults to **False**.
dolog : bool
If **True**, use a log scale to plot *var*. Defaults to **False**.
labelsize : int
Sets the font size of the labels. Defaults to 14.
title : string
Sets the plot title. Defaults to 'auto', using the variable label.
clabel : string
Set label for the color bar. Defaults to *var* and associated units.
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).
'''
from numpy import linspace, pi
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm
from matplotlib.ticker import MultipleLocator, LogLocator
# Set ax and fig based on given target.
fig, ax = set_target(target, figsize=(10.5,8), loc=loc, polar=True)
# Set function based on boolean "filled":
if filled:
contour = ax.contourf
else:
contour = ax.contour
# Get max/min if none given.
if zlim==None:
if var in self._default_zlims:
zlim = self._default_zlims[var]
else:
zlim = [self[var].min(), self[var].max()]
if dolog and zlim[0]<=0:
zlim[0] = np.min( [0.0001, zlim[1]/1000.0] )
# Create levels and set norm based on dolog.
if dolog:
levs = np.power(10, np.linspace(np.log10(zlim[0]),
np.log10(zlim[1]), nLev))
z=np.where(self[var]>zlim[0], self[var], 1.01*zlim[0])
z[z>zlim[-1]] = zlim[-1]
norm=LogNorm()
lct= LogLocator()
else:
levs = np.linspace(zlim[0], zlim[1], nLev)
z=self[var]
norm=None
lct=None
# Allow results to cross phi=360.
phi = np.concatenate( (self['lon'], [360.0]) ) * np.pi/180. - np.pi/2.
z = np.concatenate( (z, np.array([z[:,0]]).transpose() ), 1)
# Set default color tables based on variable plotted.
if ('cmap' not in kwargs) and ('colors' not in kwargs) and \
(var in self._default_cmaps):
kwargs['cmap'] = self._default_cmaps[var]
# -------Plot away------
# Contour:
cont = contour(phi, self['L'], z, levs, norm=norm, **kwargs)
# Add cbar if requested:
if add_cbar:
cbar=plt.colorbar(cont, pad=0.08, shrink=.8, ticks=lct, ax=ax)
if clabel==None:
clabel="%s ($%s$)" % (var, self[var].attrs['units'])
cbar.set_label(clabel)
else:
cbar=None
# Adjust plot appropriately.
if not Lmax:
# Default to inside ghost cells.
Lmax = self['L'][-3]
if title:
ax.set_title(title+'\n'+self.attrs['time'].isoformat(),
position=(0,1), ha='left', size=14)
_adjust_dialplot(ax, Lmax, labelsize=labelsize)
return fig, ax, cont, cbar
def add_pcolor(self, var, zlim=None, target=None, loc=111, title=None,
Lmax=None, add_cbar=False, clabel=None, dolog=False,
labelsize=14, **kwargs):
'''
Create a polar pcolor plot of variable *var* and add it to *target*.
Extra keyword arguments are handed to matplotlib's pcolor command.
Parameters
==========
var : string
The variable within the object to plot.
Returns
=======
fig : matplotlib figure object
ax : matplotlib axes object
cont : matplotlib contour object
cbar : matplotlib colorbar object
Other Parameters
================
zlim : two-element list or array
Set the color bar range. Some variables have default ranges. Others
will use max and min of *var*.
Lmax : real
Set the radial extent of the plot.
add_cbar : bool
Set whether to add a color bar or not. Defaults to **False**.
dolog : bool
If **True**, use a log scale to plot *var*. Defaults to **False**.
labelsize : int
Sets the font size of the labels. Defaults to 14.
title : string
Sets the plot title. Defaults to 'auto', using the variable label.
clabel : string
Set label for the color bar. Defaults to *var* and associated units.
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).
'''
from numpy import linspace, pi
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm
# Set ax and fig based on given target.
fig, ax = set_target(target, figsize=(10.5,8), loc=loc, polar=True)
# Get max/min if none given.
if zlim==None:
if var in self._default_zlims:
zlim = self._default_zlims[var]
else:
zlim = [self[var].min(), self[var].max()]
if dolog and zlim[0]<=0:
zlim[0] = np.min( [0.0001, zlim[1]/1000.0] )
# Logarithmic scale?
if dolog:
z=np.where(self[var]>zlim[0], self[var], 1.01*zlim[0])
norm=LogNorm()
else:
z=self[var]
norm=None
# Set up proper meshes.
nL, nPhi = len(self['L']), len(self['lon'])
dL, dPhi = self.attrs['dL'], self.attrs['dLon']*pi/180.0
phi = linspace(-1.0*dPhi/2.0, 2.*pi-dPhi/2.0, nPhi+1)-pi/2.
r = linspace(self['L'][0]-dL/2.0, self['L'][-1]+dL/2.0, nL+1)
# Set default color tables based on variable plotted.
if ('cmap' not in kwargs) and var in self._default_cmaps:
kwargs['cmap'] = self._default_cmaps[var]
# -------Plot away------
# Mesh:
pcol = ax.pcolormesh(phi, r, z, vmin=zlim[0], vmax=zlim[1],
norm=norm, **kwargs)
# Add cbar if requested:
if add_cbar:
cbar=plt.colorbar(pcol, pad=0.08, shrink=.8, ax=ax)
if clabel==None:
clabel="%s ($%s$)" % (var, self[var].attrs['units'])
cbar.set_label(clabel)
else:
cbar=None
# Adjust plot appropriately.
if not Lmax:
# Default to inside ghost cells.
Lmax = self['L'][-3]
if title:
ax.set_title(title+'\n'+self.attrs['time'].isoformat(),
position=(0,1), ha='left', size=14)
_adjust_dialplot(ax, Lmax, labelsize=labelsize)
return fig, ax, pcol, 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 add_cont(self, var, target=None, n=24, maxz=False, lines=True,
cmap=False, add_cbar=False, label=None, loc=111,
xticksize=12, yticksize=12, **kwargs):
'''
Create a polar contour of variable *var*. Plot will be either drawn
on a new matplotlib figure and axes, or you can specify a plot target
using the *target* kwarg.
Parameters
==========
var : str
The name of the variable to plot.
Returns
=======
fig : matplotlib figure object
ax : matplotlib axes object
cnt : matplotlib contour object
cb : matplotlib colorbar 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).
n : int
Set number of levels. Should be a multiple of 3 for best match
between filled and traced contours. Default is 21.
lines : bool
Add unfilled black solid/dashed contours to plot for additional
contrast. Default is **True**.
maxz : real
Set the max/min value for the color bar. Default is set by data.
cmap : str
Set the colormap. Default is to autoselect using classic IE maps.
Can be 'bwr', 'wr', or any name of a matplotlib colar map.
add_cbar : bool
Add colorbar to plot. Default is **False** which will
not add one to the plot.
label : str or None
Label to place at top of plot. Defaults to variable name.
Extra keywords are passed to the contourf routine.
'''
# Get only what we need to decrease runtime.
from math import pi
from numpy import linspace
from matplotlib.colors import Normalize
from matplotlib.ticker import MaxNLocator, MultipleLocator
from matplotlib.pyplot import clabel, colorbar
fig, ax = set_target(target, polar=True, loc=loc)
hemi = var[:2]
# Set levels and ticks:
if label==None:
label=tex_label(var)
lt_labels = ['06', label, '18', '00']
xticks = [ 0, pi/2, pi, 3*pi/2]
lct = MultipleLocator(10)
minz = self[var].min()
if minz < 0.0:
if not maxz:
maxz = max([abs(minz),self[var].max()])
crange = Normalize(vmin=-1.*maxz, vmax=maxz)
levs = linspace(-1.*maxz, maxz, n)
else:
if not maxz:
maxz = self[var].max()
crange = Normalize(vmin=0., vmax=maxz)
levs = linspace(0., maxz, n)
# Get color map if not given:
if not cmap:
if self[var].min() >= 0.0:
cmap=get_iono_cb('wr')
else:
cmap=get_iono_cb('bwr')
# Set the latitude based on hemisphere:
theta = self[hemi+'theta']
if 's_' in hemi: theta = 180-self[hemi+'theta']
# Create contour:
cnt1 = ax.contourf(self[hemi+'psi']*pi/180.0+pi/2., theta,
np.array(self[var]), levs, norm=crange, cmap=cmap)
# Set xtick label size, increase font of top label.
labels = ax.get_xticklabels()
for l in labels: l.set_size(xticksize)
labels[1].set_size(xticksize*1.25)
if lines:
nk = int(round(n/3.0))
cnt2 = ax.contour(self[hemi+'psi']*pi/180.0+pi/2., theta,
np.array(self[var]), nk, colors='k')
#clabel(cnt2,fmt='%3i',fontsize=10)
if add_cbar:
cbarticks = MaxNLocator(7)
cbar = colorbar(cnt1, ticks=cbarticks, shrink=0.75, pad=0.08)
cbar.set_label(tex_label(self[var].attrs['units']))
else:
cbar=False
ax.set_xticks(xticks)
ax.set_xticklabels(lt_labels)
ax.yaxis.set_major_locator(lct)
ax.set_ylim([0,40])
# Use text function to manually add pretty ticks.
ax.set_yticklabels('') # old ticks off.
opts = {'size':yticksize, 'rotation':-45, 'ha':'center', 'va':'center'}
for theta in [80.,70.,60.]:
txt = '{:02.0f}'.format(theta)+r'$^{\circ}$'
ax.text(pi/4., 90.-theta, txt, color='w', weight='heavy', **opts)
ax.text(pi/4., 90.-theta, txt, color='k', weight='light', **opts)
return fig, ax, cnt1, cbar
def add_separatrix(self,
target=None,
loc=111,
Lmax=None,
title=None,
**kwargs):
'''
Attempts to locate the separatrix (separator between closed and
open drift paths) by finding the minimum velocity in the domain and
tracing the path of constant potential that passes through that point.
The figure, axes, and contour object containing the separatrix line
are returned to the user.
If kwarg **target** is None (default), a new figure is
generated from scratch. If target is a matplotlib Figure
object, a new axis is created to fill that figure at subplot
location **loc**. If **target** is a matplotlib Axes object,
the plot is placed into that axis.
Four values are returned: the matplotlib Figure and Axes objects,
the matplotlib contour object, and the matplotlib colorbar object
(defaults to *False* if not used.)
Kwargs that set line characteristics behave in the typical matplotlib
manner (i.e. "colors" can be set to either a color name or a
hexidecimal specifier.)
=========== ==========================================================
Kwarg Description
----------- ----------------------------------------------------------
target Set plot destination. Defaults to new figure.
loc Set subplot location. Defaults to 111.
linewidths Set width of plotted line. Defaults to 3.0.
colors Set color of line. Defaults to 'orange'.
linestyles Set line style. Defaults to 'dashed'.
----------- ----------------------------------------------------------
'''
from numpy import linspace, pi, sqrt
import matplotlib.pyplot as plt
# Set default line behavior.
if 'linewidths' not in kwargs:
kwargs['linewidths'] = 3.0
if 'colors' not in kwargs:
kwargs['colors'] = 'orange'
if 'linestyles' not in kwargs:
kwargs['linestyles'] = 'solid'
# Set ax and fig based on given target.
fig, ax = set_target(target, figsize=(10.5, 8), loc=loc, polar=True)
doAdjust = not target == ax
# Find the stagnation point by looking for E = 0.
# Get the potential value at that spot.
# Exclude near-boundary points and limit azimuthal extent.
if 'E' not in self:
self.calc_E()
i = np.arange(self['L'].size)[self['L'] < 9.75]
j = np.arange(self['lon'].size)[(self['lon'] > 250.)
& (self['lon'] < 360.)]
pmin = self['pot'][self['E'] == self['E'][i[0]:i[-1],
j[0]:j[-1]].min()][0]
# Create a contour that only follows that potential curve.
phi = np.concatenate((self['lon'], [360.0])) * pi / 180. - pi / 2.
z = np.concatenate(
(self['pot'], np.array([self['pot'][:, 0]]).transpose()), 1)
cnt = ax.contour(phi, self['L'], z, levels=[pmin], **kwargs)
# Adjust plot appropriately.
if doAdjust:
if not Lmax:
# Default to inside ghost cells.
Lmax = self['L'][-3]
if title:
ax.set_title(title + '\n' + self.attrs['time'].isoformat(),
position=(0, 1),
ha='left',
size=14)
_adjust_dialplot(ax, Lmax, labelsize=14)
# Return bits to user.
return fig, ax, cnt
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 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 plot_df(self,
var,
time,
nlev=31,
zlim=None,
xlim=None,
ylim=None,
target=None,
loc=111,
title=None,
add_cbar=True,
clabel=None,
show_pts=False,
dolog=False,
add_body=True,
figsize=(8.34, 7),
*args,
**kwargs):
'''
Create a plot of variable *var* at a given time.
If kwarg **target** is None (default), a new figure is
generated from scratch. If target is a matplotlib Figure
object, a new axis is created to fill that figure at subplot
location **loc**. If **target** is a matplotlib Axes object,
the plot is placed into that axis.
'''
import matplotlib.pyplot as plt
from matplotlib.colors import (LogNorm, Normalize)
from matplotlib.patches import Circle
from matplotlib.ticker import (LogLocator, LogFormatter,
LogFormatterMathtext, MultipleLocator)
from matplotlib.patches import Wedge
# Grab the slice of data that we want:
if type(time) == type(self['time'][0]):
if time not in self['time']:
raise ValueError('Time not in object')
time = np.arange(self.attrs['nTime'])[self['time'] == time][0]
fig, ax = set_target(target, loc=loc, figsize=figsize)
ax.set_aspect('equal')
# Grab values from correct time/location.
x = self['Vperp[cm/s]'][time, :]
y = self['Vpar[cm/s]'][time, :]
value = self[var][time, :]
# Get max/min if none given.
if zlim == None:
zlim = [0, 0]
zlim[0] = value.min()
zlim[1] = value.max()
if dolog and zlim[0] <= 0:
zlim[0] = np.min([0.0001, zlim[1] / 1000.0])
# Create levels and set norm based on dolog.
if dolog:
levs = np.power(
10, np.linspace(np.log10(zlim[0]), np.log10(zlim[1]), nlev))
z = np.where(value > zlim[0], value, 1.01 * zlim[0])
norm = LogNorm()
ticks = LogLocator()
fmt = LogFormatterMathtext()
else:
levs = np.linspace(zlim[0], zlim[1], nlev)
z = value
norm = None
ticks = None
fmt = None
# Create contour plot.
cont=ax.tricontourf(np.asarray(x), np.asarray(y), np.asarray(z), \
np.asarray(levs), *args, norm=norm, **kwargs)
if show_pts:
ax.plot(x, y, '+w')
if xlim != None:
ax.set_xlim(xlim[0], xlim[1])
if ylim != None:
ax.set_ylim(ylim[0], ylim[1])
# Add cbar if necessary.
if add_cbar:
cbar = plt.colorbar(cont, ticks=ticks, format=fmt, pad=0.01)
if clabel == None:
clabel = "%s" % (var)
cbar.set_label(clabel)
else:
cbar = None # Need to return something, even if none.
# Set title, labels, axis ranges (use defaults where applicable.)
if title: ax.set_title(title)
#ax.set_yticks([]), ax.set_xticks([])
ax.set_xlabel('Vpar [cm/s]')
ax.set_ylabel('Vperp [cm/s]')
return fig, ax, cont, cbar
def add_cont(self, var, target=None, n=24, maxz=False, lines=True,
cmap=False, add_cbar=False, label=None, loc=111,
xticksize=12, yticksize=12, max_colat=40, **kwargs):
'''
Create a polar contour of variable *var*. Plot will be either drawn
on a new matplotlib figure and axes, or you can specify a plot target
using the *target* kwarg.
Parameters
==========
var : str
The name of the variable to plot.
Returns
=======
fig : matplotlib figure object
ax : matplotlib axes object
cnt : matplotlib contour object
cb : matplotlib colorbar 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).
n : int
Set number of levels. Should be a multiple of 3 for best match
between filled and traced contours. Default is 21.
lines : bool
Add unfilled black solid/dashed contours to plot for additional
contrast. Default is **True**.
maxz : real
Set the max/min value for the color bar. Default is set by data.
max_colat : real
Set the co-latitude range of the plot, in degrees. Defaults to 40.
cmap : str
Set the colormap. Default is to autoselect using classic IE maps.
Can be 'bwr', 'wr', or any name of a matplotlib colar map.
add_cbar : bool
Add colorbar to plot. Default is **False** which will
not add one to the plot.
label : str or None
Label to place at top of plot. Defaults to variable name.
xticksize : int
Size of longitude markers in text points. Defaults to 12.
yticksize : int
Size of latitude markers in text points. Defaults to 12.
Extra keywords are passed to the contourf routine.
'''
# Get only what we need to decrease runtime.
from math import pi
from numpy import linspace
from matplotlib.colors import Normalize
from matplotlib.ticker import MaxNLocator, MultipleLocator
from matplotlib.pyplot import clabel, colorbar
fig, ax = set_target(target, polar=True, loc=loc)
hemi = var[:2]
# user defined variables may not have hemisphere marking.
# Assume nothern hemi in those cases.
if hemi!='n_' and hemi!='s_': hemi = 'n_'
# Set levels and ticks:
if label==None:
label=tex_label(var)
lt_labels = ['06', label, '18', '00']
xticks = [ 0, pi/2, pi, 3*pi/2]
lct = MultipleLocator(10)
minz = self[var].min()
if minz < 0.0:
if not maxz:
maxz = max([abs(minz),self[var].max()])
crange = Normalize(vmin=-1.*maxz, vmax=maxz)
levs = linspace(-1.*maxz, maxz, n)
else:
if not maxz:
maxz = self[var].max()
crange = Normalize(vmin=0., vmax=maxz)
levs = linspace(0., maxz, n)
# Get color map if not given:
if not cmap:
if self[var].min() >= 0.0:
cmap=get_iono_cb('wr')
else:
cmap=get_iono_cb('bwr')
# Set the latitude based on hemisphere:
theta = self[hemi+'theta']
if 's_' in hemi: theta = 180-self[hemi+'theta']
# Create contour:
cnt1 = ax.contourf(self[hemi+'psi']*pi/180.0+pi/2., theta,
np.array(self[var]), levs, norm=crange, cmap=cmap,
**kwargs)
# Set xtick label size, increase font of top label.
labels = ax.get_xticklabels()
for l in labels: l.set_size(xticksize)
labels[1].set_size(xticksize*1.25)
if lines:
nk = int(round(n/3.0))
cnt2 = ax.contour(self[hemi+'psi']*pi/180.0+pi/2., theta,
np.array(self[var]), nk, colors='k')
#clabel(cnt2,fmt='%3i',fontsize=10)
if add_cbar:
cbarticks = MaxNLocator(7)
cbar = colorbar(cnt1, ticks=cbarticks, shrink=0.75, pad=0.08, ax=ax)
cbar.set_label(tex_label(self[var].attrs['units']))
else:
cbar=False
ax.set_xticks(xticks)
ax.set_xticklabels(lt_labels)
ax.yaxis.set_major_locator(lct)
ax.set_ylim([0,max_colat])
# Use text function to manually add pretty ticks.
ax.set_yticklabels('') # old ticks off.
opts = {'size':yticksize, 'rotation':-45, 'ha':'center', 'va':'center'}
for theta in [80.,70.,60.]:
txt = '{:02.0f}'.format(theta)+r'$^{\circ}$'
ax.text(pi/4., 90.-theta, txt, color='w', weight='heavy', **opts)
ax.text(pi/4., 90.-theta, txt, color='k', weight='light', **opts)
return fig, ax, cnt1, cbar
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 add_slice(self, var, alt, time, nlev=31, zlim=None, target=None,
loc=111, title=None, latoffset=1.05,
rlim=50., add_cbar=True, clabel=None,
show_pts=False, show_alt=True, dolog=False,
lats=[75., 60.], colats=None, figsize=(8.34,7),
*args, **kwargs):
'''
Create a plot of variable *var* at altitude *alt*.
If kwarg **target** is None (default), a new figure is
generated from scratch. If target is a matplotlib Figure
object, a new axis is created to fill that figure at subplot
location **loc**. If **target** is a matplotlib Axes object,
the plot is placed into that axis.
'''
import matplotlib.pyplot as plt
from matplotlib.colors import (LogNorm, Normalize)
from matplotlib.patches import Circle
from matplotlib.ticker import (LogLocator, LogFormatter,
LogFormatterMathtext, MultipleLocator)
# Grab the slice of data that we want:
if type(time) == type(self['time'][0]):
if time not in self['time']:
raise ValueError('Time not in object')
time = np.arange(self.attrs['nTime'])[self['time']==time][0]
fig, ax = set_target(target, loc=loc, figsize=figsize)
ax.set_aspect('equal')
# Create values in Cartesian plane.
self._get_cartXY()
# Grab values from correct time/location.
x = self._xLat[:,time, alt]
y = self._yLat[:,time, alt]
value = self[var][:, time, alt]
# Get max/min if none given.
if zlim==None:
zlim=[0,0]
zlim[0]=value.min(); zlim[1]=value.max()
if dolog and zlim[0]<=0:
zlim[0] = np.min( [0.0001, zlim[1]/1000.0] )
# Create levels and set norm based on dolog.
if dolog:
levs = np.power(10, np.linspace(np.log10(zlim[0]),
np.log10(zlim[1]), nlev))
z=np.where(value>zlim[0], value, 1.01*zlim[0])
norm=LogNorm()
ticks=LogLocator()
fmt=LogFormatterMathtext()
else:
levs = np.linspace(zlim[0], zlim[1], nlev)
z=value
norm=None
ticks=None
fmt=None
# Create contour plot.
cont=ax.tricontourf(x, y, z, levs, *args, norm=norm, **kwargs)
# Label altitude.
if show_alt:
ax.text(0.05, 0.05, r'Alt.={:.1f}$km/s$ ({:.2f}$R_E$)'.format(
self['r'][alt], self['r'][alt]/6371.0+1), size=12,
transform=ax.transAxes, bbox={'fc':'white', 'ec':'k'})
if show_pts:
ax.plot(x, y, '+w')
# Add cbar if necessary.
if add_cbar:
cbar=plt.colorbar(cont, ticks=ticks, format=fmt, pad=0.01)
if clabel==None:
clabel="%s (%s)" % (var, self[var].attrs['units'])
cbar.set_label(clabel)
else:
cbar=None # Need to return something, even if none.
# Set title, labels, axis ranges (use defaults where applicable.)
if title: ax.set_title(title)
ax.set_yticks([]), ax.set_xticks([])
ax.set_ylabel(r'Sun $\rightarrow$')
colat_lim = 90.-rlim
ax.set_xlim([-1*colat_lim, colat_lim])
ax.set_ylim([-1*colat_lim, colat_lim])
if colats:
for colat in colats:
r = latoffset*colat/np.sqrt(2)
ax.add_patch(Circle([0,0],colat,fc='none',ec='k',ls='dashed'))
ax.text(r, r, '{:.0f}'.format(colat)+r'$^{\circ}$',
rotation=315,ha='center', va='center')
else:
for lat in lats:
colat = 90 - lat
r = latoffset*colat/np.sqrt(2)
ax.add_patch(Circle([0,0],colat,fc='none',ec='k',ls='dashed'))
ax.text(r, r, '{:.0f}'.format(lat)+r'$^{\circ}$',
rotation=315,ha='center', va='center')
return fig, ax, cont, cbar
