import copy
import logging
import numpy as np
import matplotlib as mpl
from datetime import date, datetime, timezone
from matplotlib import pyplot as plt
from mpl_toolkits.axes_grid1 import make_axes_locatable
import pyspedas
from fnmatch import filter as tname_filter
from time import sleep
from pyspedas.tplot_tools import tplot_wildcard_expand, tname_byindex, get_data, var_label_panel
from pyspedas.tplot_tools import lineplot, count_traces, makegap
from pyspedas.tplot_tools import specplot, specplot_make_1d_ybins, reduce_spec_dataset
from pyspedas.tplot_tools import get_var_label_ticks
from .save_plot import save_plot
# the following improves the x-axis ticks labels
import matplotlib.units as munits
import matplotlib.dates as mdates
converter = mdates.ConciseDateConverter()
munits.registry[np.datetime64] = converter
munits.registry[date] = converter
munits.registry[datetime] = converter
def pseudovar_component_props(varname: str):
""" Return or calculate the plot properties for a single normal tplot variable
Parameters
----------
varname: str
Name of the tplot variable to inspect
Returns
-------
dict
Dictionary of plot properties for this variable::
is_spec: True if variable is a spectrogram
ymin, ymax, ylog: Y axis limits
zmin, zmax, zlog: Z axis limits (if variable is a spectrogram)
"""
output_dict = {'ymin':np.nan, 'ymax':np.nan, 'ylog':False,
'zmin':np.nan, 'zmax': np.nan, 'zlog':False,
'is_spec': False}
attrs = pyspedas.tplot_tools.data_quants[varname].attrs
plot_opts = attrs.get('plot_options')
if plot_opts is not None:
yaxis_opt = plot_opts.get('yaxis_opt')
if yaxis_opt is not None:
ylog = yaxis_opt['y_axis_type']
yrange = yaxis_opt.get('y_range')
else:
ylog = False
yrange=[None, None]
plot_extras = attrs['plot_options']['extras']
dat = get_data(varname)
if ylog is None or ylog is False or ylog == '' or ylog.lower() == 'linear':
output_dict['ylog'] = False
ylog = False
else:
output_dict['ylog'] = True
ylog = True
# ylog is now guaranteed to be boolean
if plot_extras.get('spec') is not None:
spec = plot_extras['spec']
if spec:
# Inspect a specplot variable
output_dict['is_spec']=True
if yrange is not None:
output_dict['ymin'] = yrange[0]
output_dict['ymax'] = yrange[1]
else:
# Figure out which attribute to use for bin centers (copied from specplot, maybe
# should be pulled out into a separate routine?)
var_data=dat
if len(var_data) == 3:
if hasattr(var_data, 'v'):
input_bin_centers = var_data.v
elif hasattr(var_data, 'v1'):
input_bin_centers = var_data.v1
else:
logging.warning("Multidimensional variable %s has no v or v1 attribute", varname)
elif len(var_data) == 4:
if hasattr(var_data, 'v1'):
if 'spec_dim_to_plot' in plot_extras:
if plot_extras['spec_dim_to_plot'] == 'v1':
input_bin_centers = var_data.v1
if hasattr(var_data, 'v2'):
if 'spec_dim_to_plot' in plot_extras:
if plot_extras['spec_dim_to_plot'] == 'v2':
input_bin_centers = var_data.v2
output_bin_boundaries = specplot_make_1d_ybins(None, input_bin_centers, ylog, no_regrid=True)
output_dict['ymin']=np.nanmin(output_bin_boundaries)
output_dict['ymax']=np.nanmax(output_bin_boundaries)
zaxis_options = attrs['plot_options']['zaxis_opt']
zrange = zaxis_options.get('zrange')
if zrange is not None:
output_dict['zmin'] = zrange[0]
output_dict['zmax'] = zrange[1]
else:
# Determine Z range from data array
output_dict['zmin'] = np.nanmin(dat.y)
output_dict['zmax'] = np.nanmax(dat.y)
zlog = zaxis_options.get('z_axis_type')
if zlog is None or zlog.lower() == 'linear':
output_dict['zlog'] = False
else:
output_dict['zlog'] = True
else:
# Inspect a line plot variable
output_dict['is_spec']=False
if yrange is not None:
output_dict['ymin'] = yrange[0]
output_dict['ymax'] = yrange[1]
else:
output_dict['ymin'] = np.nanmin(dat.y)
output_dict['ymax'] = np.nanmax(dat.y)
return output_dict
def gather_pseudovar_props(pseudovars:list[str]):
""" Inspect the components of a pseudovariable to determine plot limits and other settings
If any component has y or z ranges/scales set, they will be used, otherwise limits will be determined from
the data values.
If any component specifies log scaling, it will be applied to all components, otherwise linear scaling will
be used for all
Parameters
----------
pseudovars: list[str]
List of tplot composite variable components to inspect
Returns
-------
dict
A dictionary containing the plot properties to apply to each component variable::
line_ymin, line_ymax, line_ylog : Y axis limits for line plots
spec_ymin, spec_ymax, spec_ylog : Y axis limits for spectrograms
spec_zmin, spec_zmax, spec_zlog : Z axis limits for spectrograms
has_line_plots: True if any component is a line plot
has_spec_plots: True if any component is a spectrogram plot
"""
output_dict = { 'has_line_plots': False,
'has_spec_plots': False,
'line_ymin': np.nan,
'line_ymax': np.nan,
'line_ylog': False,
'spec_ymin': np.nan,
'spec_ymax': np.nan,
'spec_ylog': False,
'spec_zmin': np.nan,
'spec_zmax': np.nan,
'spec_zlog': False,
'first_spec_var': None}
for var in pseudovars:
props = pseudovar_component_props(var)
#print(f"{var}: is_spec: {props['is_spec']} ymin: {props['ymin']} ymax: {props['ymax']} ylog: {props['ylog']} zmin: {props['zmin']} zmax: {props['zmax']} zlog: {props['zlog']}")
if props['is_spec']:
output_dict['has_spec_plots'] = True
if output_dict['first_spec_var'] is None:
output_dict['first_spec_var'] = var
if props['zlog']:
output_dict['spec_zlog'] = True
if props['ylog']:
output_dict['spec_ylog'] = True
output_dict['spec_zmin']= np.nanmin([output_dict['spec_zmin'], props['zmin']])
output_dict['spec_zmax']= np.nanmax([output_dict['spec_zmax'], props['zmax']])
output_dict['spec_ymin']= np.nanmin([output_dict['spec_ymin'], props['ymin']])
output_dict['spec_ymax']= np.nanmax([output_dict['spec_ymax'], props['ymax']])
else:
output_dict['has_line_plots'] = True
if props['ylog']:
output_dict['line_ylog'] = True
output_dict['line_ymin'] = np.nanmin([output_dict['line_ymin'], props['ymin']])
output_dict['line_ymax'] = np.nanmax([output_dict['line_ymax'], props['ymax']])
return output_dict
[docs]
def tplot(variables,
trange=None,
var_label=None,
xsize=None,
ysize=None,
save_png='',
save_eps='',
save_svg='',
save_pdf='',
save_jpeg='',
dpi=None,
display=True,
fig=None,
axis=None,
running_trace_count=None,
pseudo_idx=None,
pseudo_right_axis=False,
pseudo_xaxis_options=None,
pseudo_yaxis_options=None,
pseudo_zaxis_options=None,
pseudo_line_options=None,
pseudo_extra_options=None,
show_colorbar=True,
slice=False,
return_plot_objects=False):
"""
Plot tplot variables to the display, or as saved files, using Matplotlib
Parameters
----------
variables: str or list of str, required
List of tplot variables to be plotted. Space-delimited strings may be used instead f
lists. Wildcards will be expanded.
trange: list of string or float, optional
If set, this time range will be used, temporarily overriding any previous xlim or timespan calls
var_label : str or list of str, optional
A list of variables to be displayed as values underneath the X axis major tick marks
xsize: float, optional
Plot size in the horizontal direction (in inches)
ysize: float, optional
Plot size in the vertical direction (in inches)
dpi: float, optional
The resolution of the plot in dots per inch
save_png : str, optional
A full file name and path.
If this option is set, the plot will be automatically saved to the file name provided in PNG format.
save_eps : str, optional
A full file name and path.
If this option is set, the plot will be automatically saved to the file name provided in EPS format.
save_jpeg : str, optional
A full file name and path.
If this option is set, the plot will be automatically saved to the file name provided in JPEG format.
save_pdf : str, optional
A full file name and path.
If this option is set, the plot will be automatically saved to the file name provided in PDF format.
save_svg : str, optional
A full file name and path.
If this option is set, the plot will be automatically saved to the file name provided in SVG format.
dpi: float, optional
The resolution of the plot in dots per inch
display: bool, optional
If True, then this function will display the plotted tplot variables. Use False to suppress display (for example, if
saving to a file, or returning plot objects to be displayed later). Default: True
fig: Matplotlib figure object
Use an existing figure to plot in (mainly for recursive calls to render composite variables)
axis: Matplotlib axes object
Use an existing set of axes to plot on (mainly for recursive calls to render composite variables)
running_trace_count: int, optional
In recursive calls for rendering composite variables, the index of the trace currently being rendered.
pseudo_idx: int, optional
In recursive calls for rendering composite variables, the index of the variable component currently being rendered.
pseudo_right_axis: bool, optional
In recursive calls for rendering composite variables, a flag to indicate the Y scale should be placed on
the right Y axis.
pseudo_xaxis_options: dict, optional
In recursice calls for rendering composite variables, X axis options inherited from the parent variable
pseudo_yaxis_options: dict, optional
In recursive calls for rendering composite variables, Y axis options inherited from the parent variable
pseudo_zaxis_options: dict, optional
In recursive calls for rendering composite variables, Z axis options inherited from the parent variable
pseudo_line_options: dict, optional
In recursive calls for rendering composite variables, line options inherited from the parent variable
pseudo_extra_options: dict, optional
In recursive calls for rendering composite variables, extra options inherited from the parent variable
show_colorbar: bool, optional
Show a colorbar showing the Z scale for spectrogram plots.
slice: bool, optional
If True, show an interactive window with a plot of Z versus Y values for the X axis (time) value under the cursor. Default: False
return_plot_objects: bool, optional
If true, returns the matplotlib fig and axes objects for further manipulation. Default: False
Returns
-------
Any
Returns matplotlib fig and axes objects, if return_plot_objects==True
Examples
--------
>>> # Plot a single line plot
>>> import pyspedas
>>> x_data = [2,3,4,5,6]
>>> y_data = [1,2,3,4,5]
>>> pyspedas.store_data("Variable1", data={'x':x_data, 'y':y_data})
>>> pyspedas.tplot("Variable1")
>>> # Plot two variables
>>> x_data = [1,2,3,4,5]
>>> y_data = [[1,5],[2,4],[3,3],[4,2],[5,1]]
>>> pyspedas.store_data("Variable2", data={'x':x_data, 'y':y_data})
>>> pyspedas.tplot(["Variable1", "Variable2"])
>>> # Plot two plots, adding a third variable's values as annotations at X axis major tick marks
>>> x_data = [1,2,3]
>>> y_data = [ [1,2,3] , [4,5,6], [7,8,9] ]
>>> v_data = [1,2,3]
>>> pyspedas.store_data("Variable3", data={'x':x_data, 'y':y_data, 'v':v_data})
>>> pyspedas.options("Variable3", 'spec', 1)
>>> pyspedas.tplot(["Variable2", "Variable3"], var_label='Variable1')
"""
# This call resolves wildcard patterns and converts integers to variable names
variables = tplot_wildcard_expand(variables)
if len(variables) == 0:
logging.warning("tplot: No matching tplot names were found")
return
varlabel_style = pyspedas.tplot_tools.tplot_opt_glob.get('varlabel_style')
if varlabel_style is None or varlabel_style.lower() == 'extra_axes':
num_panels = len(variables)
panel_sizes = [1]*num_panels
else: # varlabel_style 'extra_panel'
num_panels = len(variables) + 1
panel_sizes = [1] * len(variables) + [0.1 * (len(var_label) + 2)]
# support for the panel_size option
for var_idx, variable in enumerate(variables):
if pyspedas.tplot_tools.data_quants.get(variable) is None:
continue
panel_size = pyspedas.tplot_tools.data_quants[variable].attrs['plot_options']['extras'].get('panel_size')
if panel_size is not None:
panel_sizes[var_idx] = panel_size
if xsize is None:
xsize = pyspedas.tplot_tools.tplot_opt_glob.get('xsize')
if xsize is None:
xsize = 12
if ysize is None:
ysize = pyspedas.tplot_tools.tplot_opt_glob.get('ysize')
if ysize is None:
if num_panels > 4:
ysize = 8
else:
# This was previously set to 5, which resulted in a non-monotonic progression of panel heights
default_y_sizes = [5, 5, 6, 7, 8]
ysize = default_y_sizes[num_panels]
# The logic here for handling the 'right_axis' option is pretty convoluted, and makes a number of assumptions
# that may not be warranted: mainly that right_axis will only be set on pseudovariables, and that if
# set, the first sub-variable gets the left axis and all other sub-variables get a newly twinx-ed right axis.
# "right axis" implies there will only be at most two Y axes. What if more scales are needed?
# There also seems to be some conflation of the set of axes for the whole stack of plots, versus
# the single (left or possibly right) axis for the variable currently being rendered.
#
# This whole concept is kind of a mess at the moment. For now, we'll make it work for the
# most likely use case, plotting a single spectrum variable followed by a single line variable
# (for example, an energy spectrum plus spacecraft potential). JWL 2024-03-26
if fig is None and axis is None:
fig, axes = plt.subplots(nrows=num_panels, sharex=True, gridspec_kw={'height_ratios': panel_sizes}, layout='constrained')
fig.set_size_inches(xsize, ysize)
plot_title = pyspedas.tplot_tools.tplot_opt_glob['title_text']
fig.suptitle(plot_title)
if (plot_title is not None) and plot_title != '':
if 'title_size' in pyspedas.tplot_tools.tplot_opt_glob:
title_size = pyspedas.tplot_tools.tplot_opt_glob['title_size']
fig.suptitle(plot_title, fontsize=title_size)
else:
fig.suptitle(plot_title)
# support for matplotlib styles
style = pyspedas.tplot_tools.tplot_opt_glob.get('style')
if style is not None:
plt.style.use(style)
else:
# fig and axis have been passed as parameters, most likely a recursive tplot call to render
# a pseudovariable
if pseudo_idx == 0 or pseudo_right_axis == False:
# setting up first axis
axes = axis
elif pseudo_idx > 0 and pseudo_right_axis:
# generate and use the right axis? probably still wrong...
axes = axis.twinx()
# support for matplotlib styles -- shouldn't be needed if processing pseudovar components?
style = pyspedas.tplot_tools.tplot_opt_glob.get('style')
if style is not None:
plt.style.use(style)
axis_font_size = pyspedas.tplot_tools.tplot_opt_glob.get('axis_font_size')
colorbars = {}
for idx, variable in enumerate(variables):
var_data_org = get_data(variable, dt=True)
var_metadata = get_data(variable, metadata=True)
#Check for a 3d variable, call reduce_spec_dataset
if hasattr(var_data_org, 'v1') and hasattr(var_data_org, 'v2'):
temp_dq = reduce_spec_dataset(name=variable)
var_data_org = get_data(variable, dt=True, data_quant_in=temp_dq)
if var_data_org is None:
logging.info('Variable not found: ' + variable)
continue
var_data = copy.deepcopy(var_data_org)
# plt.subplots returns a list of axes for multiple panels
# but only a single axis for a single panel
if num_panels == 1:
this_axis = axes
else:
this_axis = axes[idx]
# we need to track the variable name in the axis object
# for spectrogram slices
this_axis.var_name = variable
pseudo_var = False
overplots = None
spec = False
var_quants = pyspedas.tplot_tools.data_quants[variable]
if not isinstance(var_quants, dict):
overplots = var_quants.attrs['plot_options'].get('overplots_mpl')
if overplots is not None and len(overplots) > 0:
pseudo_var = True
# deal with pseudo-variables first
if isinstance(var_data, list) or isinstance(var_data, str) or pseudo_var:
# this is a pseudo variable
if isinstance(var_data, str):
var_data = var_data.split(' ')
if pseudo_var:
pseudo_vars = overplots
else:
pseudo_vars = var_data
# pseudo variable metadata should override the metadata
# for individual variables
xaxis_options = None
yaxis_options = None
zaxis_options = None
line_opts = None
plot_extras = None
if pseudo_var:
plot_extras = var_quants.attrs['plot_options']['extras']
if plot_extras.get('spec') is not None:
spec = True
if plot_extras.get('right_axis') is not None:
if plot_extras.get('right_axis'):
pseudo_right_axis = True
if pseudo_right_axis or spec:
plot_extras = None
else:
yaxis_options = var_quants.attrs['plot_options']['yaxis_opt']
zaxis_options = var_quants.attrs['plot_options']['zaxis_opt']
line_opts = var_quants.attrs['plot_options']['line_opt']
traces_processed = 0
pseudovar_props = gather_pseudovar_props(pseudo_vars)
#
# Determine which Y axis options need to be changed to accommodate multiple component variables
# There might be line variables and spectra, each with their own overall yscale and yrange.
# If both are present, the spectra will take priority.
# If explicit y or z scales or ranges are set on the pseudovariable, do not override.
if pseudovar_props['has_spec_plots']:
if pseudovar_props['spec_ylog']:
new_yscale = 'log'
else:
new_yscale = 'linear'
new_yr = [pseudovar_props['spec_ymin'], pseudovar_props['spec_ymax']]
else:
if pseudovar_props['line_ylog']:
new_yscale = 'log'
else:
new_yscale = 'linear'
new_yr = [pseudovar_props['line_ymin'], pseudovar_props['line_ymax']]
override_yopts = {}
if pseudo_right_axis:
# Who knows? Let the component variable Y-axis fight it out...
override_yopts = {}
elif yaxis_options is None:
override_yopts = {'y_range':new_yr, 'y_range_user':True,'y_axis_style':new_yscale}
else:
if yaxis_options.get('y_range') is None:
override_yopts['y_range'] = new_yr
override_yopts['y_range_user'] = True
if yaxis_options.get('y_axis_style') is None:
override_yopts['y_axis_style'] = new_yscale
if yaxis_options is not None:
yaxis_options = yaxis_options | override_yopts
else:
yaxis_options = override_yopts
# Determine which Z axis options need to be changed to accommodate multiple component variables
#
# We only care about specplots here.
new_zr = [ None, pseudovar_props['spec_zmax']]
if pseudovar_props['spec_zlog']:
new_zscale='log'
else:
new_zscale='linear'
override_zopts = {}
if zaxis_options is None:
override_zopts = {'z_range':new_zr, 'z_axis_style':new_zscale}
else:
if zaxis_options.get('z_range') is None:
override_zopts['z_range'] = new_zr
if zaxis_options.get('z_axis_style') is None:
override_zopts['z_axis_style'] = new_zscale
if zaxis_options is not None:
zaxis_options = zaxis_options | override_zopts
else:
zaxis_options = override_zopts
for pseudo_idx, var in enumerate(pseudo_vars):
# We're plotting a pseudovariable. Iterate over the sub-variables, keeping track of how many
# traces have been plotted so far, so we can correctly match option values to traces. The pseudovariable
# y_axis, z_axis, line and extra options are passed as parameters so they can be merged with the
# sub-variable options, with any pseudovar options overriding the sub-variable options.
trace_count_thisvar = count_traces(var)
if var == pseudovar_props['first_spec_var']:
pseudo_show_colorbar=True
else:
pseudo_show_colorbar=False
tplot(var,
trange=trange,
return_plot_objects=return_plot_objects,
xsize=xsize, ysize=ysize,
fig=fig, axis=this_axis, display=False,
running_trace_count=traces_processed,
pseudo_idx=pseudo_idx,
pseudo_xaxis_options=xaxis_options, pseudo_yaxis_options=yaxis_options, pseudo_zaxis_options=zaxis_options,
pseudo_line_options=line_opts, pseudo_extra_options=plot_extras,
pseudo_right_axis=pseudo_right_axis,
show_colorbar=pseudo_show_colorbar)
traces_processed += trace_count_thisvar
continue
#if data_gap is an option for this variable, or if it's a add
#gaps here; an individual gap setting should override the
#global setting
plot_extras = var_quants.attrs['plot_options']['extras']
if plot_extras.get('data_gap') is not None and plot_extras.get('data_gap') > 0:
var_data = makegap(var_data, dt = plot_extras.get('data_gap'))
else:
if pyspedas.tplot_tools.tplot_opt_glob['data_gap'] is not None and pyspedas.tplot_tools.tplot_opt_glob['data_gap'] > 0:
var_data = makegap(var_data, dt = pyspedas.tplot_tools.tplot_opt_glob['data_gap'])
# set the x-axis range, if it was set with xlim or tlimit or the trange parameter
if trange is None and pyspedas.tplot_tools.tplot_opt_glob.get('x_range') is None:
var_data_times = var_data.times
time_idxs = np.arange(len(var_data_times))
else:
if trange is not None:
if len(trange) != 2:
logging.error('Invalid trange setting: must be a 2-element list or array')
return
if isinstance(trange[0], str):
x_range = pyspedas.tplot_tools.time_double(trange) # seconds since epoch
x_range_start = x_range[0]
x_range_stop = x_range[1]
else:
x_range = pyspedas.tplot_tools.tplot_opt_glob['x_range'] # Seconds since epoch
x_range_start = x_range[0]
x_range_stop = x_range[1]
# Check for NaN or inf in x_range
if not np.isfinite(x_range_start):
logging.warning('tplot: x_range start is not finite, replacing with 0')
x_range_start = 0
if not np.isfinite(x_range_stop):
logging.warning('tplot: x_range end is not finite, replacing with 0')
x_range_stop = 0
# Convert to np.datetime64 with nanosecond precision
x_range = np.array(np.array([x_range_start*1e9, x_range_stop*1e9]),dtype='datetime64[ns]')
this_axis.set_xlim(x_range)
time_idxs = np.argwhere((var_data.times >= x_range[0]) & (var_data.times <= x_range[1])).flatten()
if len(time_idxs) == 0:
logging.info('No data found in the time range: ' + variable)
continue
var_data_times = var_data.times[time_idxs]
var_times = var_data_times
# set some more plot options
xaxis_options = var_quants.attrs['plot_options']['xaxis_opt']
if pseudo_xaxis_options is not None and len(pseudo_xaxis_options) > 0:
merged_xaxis_options = xaxis_options | pseudo_xaxis_options
xaxis_options = merged_xaxis_options
yaxis_options = var_quants.attrs['plot_options']['yaxis_opt']
if pseudo_yaxis_options is not None and len(pseudo_yaxis_options) > 0:
merged_yaxis_options = yaxis_options | pseudo_yaxis_options
yaxis_options = merged_yaxis_options
zaxis_options = var_quants.attrs['plot_options']['zaxis_opt']
if pseudo_zaxis_options is not None and len(pseudo_zaxis_options) > 0:
merged_zaxis_options = zaxis_options | pseudo_zaxis_options
zaxis_options = merged_zaxis_options
line_opts = var_quants.attrs['plot_options']['line_opt']
if pseudo_line_options is not None and len(pseudo_line_options) > 0:
merged_line_opts = line_opts | pseudo_line_options
line_opts = merged_line_opts
if line_opts is not None:
if 'name' in line_opts:
this_axis.set_title(line_opts['name'])
elif 'title' in line_opts:
this_axis.set_title(line_opts['title'])
plot_extras = var_quants.attrs['plot_options']['extras']
if pseudo_extra_options is not None and len(pseudo_extra_options) > 0:
merged_plot_extras = plot_extras | pseudo_extra_options
plot_extras = merged_plot_extras
xtitle = None
if xaxis_options.get('axis_label') is not None:
xtitle = xaxis_options['axis_label']
xsubtitle = ''
if xaxis_options.get('axis_subtitle') is not None:
xsubtitle = xaxis_options['axis_subtitle']
if style is None:
xtitle_color = 'black'
else:
xtitle_color = None
if xaxis_options.get('axis_color') is not None:
xtitle_color = xaxis_options['axis_color']
ylog = yaxis_options['y_axis_type']
if ylog == 'log':
this_axis.set_yscale('log')
else:
this_axis.set_yscale('linear')
ytitle = yaxis_options['axis_label']
if ytitle == '':
ytitle = variable
ysubtitle = ''
if yaxis_options.get('axis_subtitle') is not None:
ysubtitle = yaxis_options['axis_subtitle']
# replace some common superscripts
ysubtitle = replace_common_exp(ysubtitle)
if axis_font_size is not None:
this_axis.tick_params(axis='x', labelsize=axis_font_size)
this_axis.tick_params(axis='y', labelsize=axis_font_size)
char_size = pyspedas.tplot_tools.tplot_opt_glob.get('charsize')
if char_size is None:
char_size = 12
if plot_extras.get('char_size') is not None:
char_size = plot_extras['char_size']
user_set_yrange = yaxis_options.get('y_range_user')
if user_set_yrange is not None:
# the user has set the yrange manually
yrange = yaxis_options['y_range']
if not np.isfinite(yrange[0]):
yrange[0] = None
if not np.isfinite(yrange[1]):
yrange[1] = None
this_axis.set_ylim(yrange)
ymajor_ticks = yaxis_options.get('y_major_ticks')
if ymajor_ticks is not None:
this_axis.set_yticks(ymajor_ticks)
yminor_tick_interval = yaxis_options.get('y_minor_tick_interval')
if yminor_tick_interval is not None and ylog != 'log':
this_axis.yaxis.set_minor_locator(plt.MultipleLocator(yminor_tick_interval))
if style is None:
ytitle_color = 'black'
else:
ytitle_color = None
if yaxis_options.get('axis_color') is not None:
ytitle_color = yaxis_options['axis_color']
if xtitle is not None and xtitle != '':
if xtitle_color is not None:
this_axis.set_xlabel(xtitle + '\n' + xsubtitle, fontsize=char_size, color=xtitle_color)
else:
this_axis.set_xlabel(xtitle + '\n' + xsubtitle, fontsize=char_size)
if ytitle_color is not None:
this_axis.set_ylabel(ytitle + '\n' + ysubtitle, fontsize=char_size, color=ytitle_color)
else:
this_axis.set_ylabel(ytitle + '\n' + ysubtitle, fontsize=char_size)
border = True
if plot_extras.get('border') is not None:
border = plot_extras['border']
if border == False:
this_axis.axis('off')
# axis tick options
if plot_extras.get('xtickcolor') is not None:
this_axis.tick_params(axis='x', color=plot_extras.get('xtickcolor'))
if plot_extras.get('ytickcolor') is not None:
this_axis.tick_params(axis='y', color=plot_extras.get('ytickcolor'))
if plot_extras.get('xtick_direction') is not None:
this_axis.tick_params(axis='x', direction=plot_extras.get('xtick_direction'))
if plot_extras.get('ytick_direction') is not None:
this_axis.tick_params(axis='y', direction=plot_extras.get('ytick_direction'))
if plot_extras.get('xtick_length') is not None:
this_axis.tick_params(axis='x', length=plot_extras.get('xtick_length'))
if plot_extras.get('ytick_length') is not None:
this_axis.tick_params(axis='y', length=plot_extras.get('ytick_length'))
if plot_extras.get('xtick_width') is not None:
this_axis.tick_params(axis='x', width=plot_extras.get('xtick_width'))
if plot_extras.get('ytick_width') is not None:
this_axis.tick_params(axis='y', width=plot_extras.get('ytick_width'))
if plot_extras.get('xtick_labelcolor') is not None:
this_axis.tick_params(axis='x', labelcolor=plot_extras.get('xtick_labelcolor'))
if plot_extras.get('ytick_labelcolor') is not None:
this_axis.tick_params(axis='y', labelcolor=plot_extras.get('ytick_labelcolor'))
# determine if this is a line plot or a spectrogram
spec = False
if plot_extras.get('spec') is not None:
spec = plot_extras['spec']
if spec:
# create spectrogram plots
plot_created = specplot(var_data, var_times, this_axis, yaxis_options, zaxis_options, plot_extras, colorbars, axis_font_size, fig, variable, time_idxs=time_idxs, style=style)
if not plot_created:
continue
else:
# create line plots
plot_created = lineplot(var_data, var_times, this_axis, line_opts, yaxis_options, plot_extras, running_trace_count=running_trace_count, time_idxs=time_idxs, style=style, var_metadata=var_metadata)
if not plot_created:
continue
# apply any vertical/horizontal bars
if pyspedas.tplot_tools.data_quants[variable].attrs['plot_options'].get('time_bar') is not None:
time_bars = pyspedas.tplot_tools.data_quants[variable].attrs['plot_options']['time_bar']
for time_bar in time_bars:
# vertical bars
if time_bar['dimension'] == 'height':
this_axis.axvline(x=datetime.fromtimestamp(time_bar['location'], tz=timezone.utc),
color=np.array(time_bar.get('line_color'))/256.0, lw=time_bar.get('line_width'),
linestyle=time_bar.get('line_dash'))
# horizontal bars
if time_bar['dimension'] == 'width':
this_axis.axhline(y=time_bar['location'], color=np.array(time_bar.get('line_color'))/256.0,
lw=time_bar.get('line_width'),
linestyle=time_bar.get('line_dash'))
# highlight time intervals
if pyspedas.tplot_tools.data_quants[variable].attrs['plot_options'].get('highlight_intervals') is not None:
highlight_intervals = pyspedas.tplot_tools.data_quants[variable].attrs['plot_options']['highlight_intervals']
for highlight_interval in highlight_intervals:
hightlight_opts = copy.deepcopy(highlight_interval)
del hightlight_opts['location']
if highlight_interval['edgecolor'] is not None or highlight_interval['facecolor'] is not None:
del hightlight_opts['color']
this_axis.axvspan(mdates.date2num(datetime.fromtimestamp(highlight_interval['location'][0], timezone.utc)),
mdates.date2num(datetime.fromtimestamp(highlight_interval['location'][1], timezone.utc)),
**hightlight_opts)
# add annotations
if pyspedas.tplot_tools.data_quants[variable].attrs['plot_options']['extras'].get('annotations') is not None:
annotations = pyspedas.tplot_tools.data_quants[variable].attrs['plot_options']['extras']['annotations']
for annotation in annotations:
this_axis.annotate(annotation['text'], annotation['position'],
xycoords=annotation['xycoords'],
fontsize=annotation['fontsize'],
alpha=annotation['alpha'],
fontfamily=annotation['fontfamily'],
fontvariant=annotation['fontvariant'],
fontstyle=annotation['fontstyle'],
fontstretch=annotation['fontstretch'],
fontweight=annotation['fontweight'],
rotation=annotation['rotation'],
color=annotation['color'])
# apply any addition x-axes (or panel) specified by the var_label keyword
if var_label is not None:
if varlabel_style is None or varlabel_style.lower() == 'extra_axes':
varlabels_extra_axes(num_panels, this_axis, var_label, axis_font_size, plot_extras=plot_extras)
else:
var_label_panel(variables, var_label, axes, axis_font_size)
# add the color bars to any spectra
for idx, variable in enumerate(variables):
if pyspedas.tplot_tools.data_quants.get(variable) is None:
continue
plot_extras = pyspedas.tplot_tools.data_quants[variable].attrs['plot_options']['extras']
zaxis_options = pyspedas.tplot_tools.data_quants[variable].attrs['plot_options']['zaxis_opt']
if plot_extras.get('spec') is not None:
spec = plot_extras['spec']
else:
spec = False
if spec and show_colorbar:
if colorbars.get(variable) is None:
continue
if num_panels == 1:
this_axis = axes
else:
this_axis = axes[idx]
colorbar = fig.colorbar(colorbars[variable]['im'], ax=this_axis)
if style is None:
ztitle_color = 'black'
else:
ztitle_color = None
if zaxis_options is None:
continue
if zaxis_options.get('axis_color') is not None:
ztitle_color = zaxis_options['axis_color']
ztitle_text = colorbars[variable]['ztitle']
zsubtitle_text = colorbars[variable]['zsubtitle']
# replace some common superscripts
ztitle_text = replace_common_exp(ztitle_text)
zsubtitle_text = replace_common_exp(zsubtitle_text)
if ztitle_color is not None:
colorbar.set_label(ztitle_text + '\n ' + zsubtitle_text,
color=ztitle_color, fontsize=char_size)
else:
colorbar.set_label(ztitle_text + '\n ' + zsubtitle_text,
fontsize=char_size)
# plt.tight_layout()
fig.canvas.draw()
save_plot(save_png=save_png, save_eps=save_eps, save_jpeg=save_jpeg, save_pdf=save_pdf, save_svg=save_svg, dpi=dpi)
if slice:
slice_fig, slice_axes = plt.subplots(nrows=1)
slice_plot, = slice_axes.plot([0], [0])
mouse_event_func = lambda event: mouse_move_slice(event, slice_axes, slice_plot)
cid = fig.canvas.mpl_connect('motion_notify_event', mouse_event_func)
if display:
plt.show()
if return_plot_objects:
return fig, axes
def varlabels_extra_axes(num_panels, this_axis, var_label, axis_font_size, plot_extras):
# apply any addition x-axes specified by the var_label keyword
if var_label is not None:
if not isinstance(var_label, list):
var_label = [var_label]
char_size = pyspedas.tplot_tools.tplot_opt_glob.get('charsize')
if char_size is None:
char_size = 12
if plot_extras.get('char_size') is not None:
char_size = plot_extras['char_size']
axis_delta = 0.0
for label in var_label:
if isinstance(label, int):
label = tname_byindex(label)
label_data = get_data(label, xarray=True, dt=True)
if label_data is None:
logging.info('Variable not found: ' + label)
continue
if len(label_data.values.shape) != 1:
logging.info(
label + ' specified as a vector; var_label only supports scalars. Try splitting the vector into seperate tplot variables.')
continue
# set up the new x-axis
axis_delta = axis_delta - num_panels * 0.1
new_xaxis = this_axis.secondary_xaxis(axis_delta)
if axis_font_size is not None:
new_xaxis.tick_params(axis='x', labelsize=axis_font_size)
new_xaxis.tick_params(axis='y', labelsize=axis_font_size)
xaxis_ticks = this_axis.get_xticks().tolist()
xaxis_ticks_dt = [np.datetime64(mpl.dates.num2date(tick_val).replace(tzinfo=None).isoformat(), 'ns') for
tick_val in xaxis_ticks]
# xaxis_ticks_unix = [tick_val.timestamp() for tick_val in xaxis_ticks_dt]
xaxis_labels = get_var_label_ticks(label_data, xaxis_ticks_dt)
new_xaxis.set_xticks(xaxis_ticks_dt)
new_xaxis.set_xticklabels(xaxis_labels)
ytitle = pyspedas.tplot_tools.data_quants[label].attrs['plot_options']['yaxis_opt']['axis_label']
new_xaxis.set_xlabel(ytitle, fontsize=char_size)
# fig.subplots_adjust(bottom=0.05+len(var_label)*0.1)
def mouse_move_slice(event, slice_axes, slice_plot):
"""
This function is called when the mouse moves over an axis
and the slice keyword is set to True; for spectra figures, it
updates the slice plot based on the mouse location
"""
if event.inaxes is None:
return
# check for a spectrogram
try:
data = get_data(event.inaxes.var_name)
except AttributeError:
return
if data is None:
return
if len(data) != 3:
return
slice_time = mdates.num2date(event.xdata).timestamp()
idx = np.abs(data.times-slice_time).argmin()
if len(data.v.shape) > 1:
# time varying y-axis
vdata = data.v[idx, :]
else:
vdata = data.v
xaxis_options = pyspedas.tplot_tools.data_quants[event.inaxes.var_name].attrs['plot_options']['xaxis_opt']
yaxis_options = pyspedas.tplot_tools.data_quants[event.inaxes.var_name].attrs['plot_options']['yaxis_opt']
zaxis_options = pyspedas.tplot_tools.data_quants[event.inaxes.var_name].attrs['plot_options']['zaxis_opt']
yrange = yaxis_options.get('y_range')
if yrange is None:
yrange = [np.nanmin(vdata), np.nanmax(vdata)]
zrange = zaxis_options.get('z_range')
if zrange is None:
zrange = [np.nanmin(data.y), np.nanmax(data.y)]
y_label = zaxis_options.get('axis_label')
if y_label is not None:
slice_axes.set_ylabel(y_label)
title = datetime.fromtimestamp(data.times[idx], timezone.utc).strftime('%Y-%m-%d %H:%M:%S.%f')
x_label = yaxis_options.get('axis_label')
if x_label is not None:
title = x_label + ' (' + title + ')'
slice_axes.set_title(title)
x_subtitle = yaxis_options.get('axis_subtitle')
if x_subtitle is not None:
slice_axes.set_xlabel(x_subtitle)
slice_yaxis_opt = pyspedas.tplot_tools.data_quants[event.inaxes.var_name].attrs['plot_options'].get('slice_yaxis_opt')
xscale = None
yscale = None
if slice_yaxis_opt is not None:
xscale = slice_yaxis_opt.get('xi_axis_type')
yscale = slice_yaxis_opt.get('yi_axis_type')
if yscale is None:
# if the user didn't explicitly set the ylog_slice option,
# use the option from the plot
yscale = zaxis_options.get('z_axis_type')
if yscale is None:
yscale = 'linear'
if xscale is None:
# if the user didn't explicitly set the xlog_slice option,
# use the option from the plot
xscale = yaxis_options.get('y_axis_type')
if xscale is None:
xscale = 'linear'
if yscale == 'log' and zrange[0] == 0.0:
zrange[0] = np.nanmin(data.y[idx, :])
slice_plot.set_data(vdata, data.y[idx, :])
slice_axes.set_ylim(zrange)
slice_axes.set_xlim(yrange)
slice_axes.set_xscale(xscale)
slice_axes.set_yscale(yscale)
try:
plt.draw()
except ValueError:
return
sleep(0.01)
def replace_common_exp(title):
if hasattr(title, 'decode'):
title = title.decode('utf-8')
if '$' in title:
return title
if '^' not in title:
return title
exp = False
title_out = ''
for char in title:
if char == '^':
exp = True
title_out += '$^{'
continue
else:
if exp:
if not char.isalnum():
title_out += '}$' + char
exp = False
continue
title_out += char
if exp:
title_out += '}$'
return title_out