import numpy as np
import pyspedas
import logging
[docs]
def lineplot(var_data,
var_times,
this_axis,
line_opts,
yaxis_options,
plot_extras,
running_trace_count=None,
time_idxs=None,
style=None,
var_metadata=None):
"""
Generate a matplotlib line plot from a tplot variable
Parameters
----------
var_data: dict
The data to be plotted (may have multiple traces)
var_times:
Array of datetime objects to use for x axis
this_axis:
The current axis (plot panel) we're working with
line_opts: dict
A dictionary of line options
yaxis_options: dict
A dictionary of y axis options
plot_extras: dict
A dictionary of 'extra' options (colors, etc)
running_trace_count:
If not Null, an integer representing the number of traces already processed in this pseudovariable. Defaults to None.
time_idxs: np.ndarray
If provided, an integer array specifying the subset of time indices to be plotted. Defaults to None.
style:
A matplotlib style to be used in the plot. Defaults to None.
var_metadata: dict
The metadata dictionary associated with this tplot variable (used as a fallback for trace labels). Defaults to None.
Returns
-------
True
"""
alpha = plot_extras.get('alpha')
if len(var_data.y.shape) == 1:
num_lines = 1
else:
num_lines = var_data.y.shape[1]
is_errorbar_plot = False
if 'dy' in var_data._fields:
is_errorbar_plot = True
if yaxis_options.get('legend_names') is not None:
labels = yaxis_options['legend_names']
labels = get_trace_options(labels, running_trace_count, num_lines)
if labels[0] is None:
labels = None
else:
labels = None
if var_metadata.get('CDF') is not None:
labels = var_metadata['CDF'].get('LABELS')
legend_location = yaxis_options.get('legend_location')
bbox_to_anchor = None
if legend_location is not None:
if legend_location == 'spedas':
# the spedas legend puts the legend on the outside of the panel
# to the right of the panel (just like in IDL)
legend_location = 'center left'
bbox_to_anchor = (1.04, 0.5)
else:
legend_location = 'upper right'
legend_size = yaxis_options.get('legend_size')
legend_shadow = yaxis_options.get('legend_shadow')
legend_title = yaxis_options.get('legend_title')
legend_titlesize = yaxis_options.get('legend_titlesize')
legend_color = yaxis_options.get('legend_color')
legend_markerfirst = yaxis_options.get('legend_markerfirst')
legend_markerscale = yaxis_options.get('legend_markerscale')
legend_linewidth = yaxis_options.get('legend_linewidth')
legend_edgecolor = yaxis_options.get('legend_edgecolor')
legend_facecolor = yaxis_options.get('legend_facecolor')
legend_frameon = yaxis_options.get('legend_frameon')
legend_ncols = yaxis_options.get('legend_ncols')
if legend_ncols is None:
legend_ncols = 1
if legend_linewidth is None:
legend_linewidth = 4
if legend_size is None:
legend_size = pyspedas.tplot_tools.tplot_opt_glob.get('charsize')
markers = None
if line_opts.get('marker') is not None:
markers = line_opts['marker']
markers = get_trace_options(markers, running_trace_count, num_lines, repeat=True)
colors = None
if plot_extras.get('line_color') is not None:
colors = plot_extras['line_color']
else:
if style is None:
if num_lines == 1:
colors = ['k']
elif num_lines == 2:
colors = ['r', 'g']
elif num_lines == 3:
colors = ['b', 'g', 'r']
elif num_lines == 4:
colors = ['b', 'g', 'r', 'k']
else:
colors = ['k', 'C2', 'C3', 'C4', 'C5', 'C6', 'C7', 'C8', 'C9']
colors = get_trace_options(colors, running_trace_count, num_lines, repeat=True)
# line thickness
if line_opts.get('line_width') is not None:
thick = line_opts['line_width']
else:
thick = [0.5]
thick = get_trace_options(thick, running_trace_count, num_lines, repeat=True)
# line style
if line_opts.get('line_style_name') is not None:
line_style_user = line_opts['line_style_name']
# line_style_user should already be a list
# handle legacy values
line_style = []
for linestyle in line_style_user:
if linestyle == 'solid_line':
line_style.append('solid')
elif linestyle == 'dot':
line_style.append('dotted')
elif linestyle == 'dash':
line_style.append('dashed')
elif linestyle == 'dash_dot':
line_style.append('dashdot')
else:
line_style.append(linestyle)
else:
line_style = ['solid']
line_style = get_trace_options(line_style, running_trace_count, num_lines, repeat=True)
symbols = False
if line_opts.get('symbols') is not None:
if line_opts['symbols']:
symbols = True
# create the plot
line_options = {'alpha': alpha}
marker_every = None
if line_opts.get('markevery') is not None:
marker_every = line_opts['markevery']
marker_every = get_trace_options(marker_every, running_trace_count, num_lines, repeat=True)
marker_sizes = None
if line_opts.get('marker_size') is not None:
marker_sizes = line_opts['marker_size']
marker_sizes = get_trace_options(marker_sizes, running_trace_count, num_lines, repeat=True)
# check for error data first
if is_errorbar_plot:
# error data provided
line_options['yerr'] = var_data.dy[time_idxs]
plotter = this_axis.errorbar
if line_opts.get('ecolor') is not None:
line_options['ecolor'] = line_opts['ecolor']
if line_opts.get('elinewidth') is not None:
line_options['elinewidth'] = line_opts['elinewidth']
if line_opts.get('errorevery') is not None:
line_options['errorevery'] = line_opts['errorevery']
if line_opts.get('capsize') is not None:
line_options['capsize'] = line_opts['capsize']
else:
# no error data provided
plotter = this_axis.plot
# Note: to turn off connecting lines in an error bar plot, do not use the
# 'symbols' option. Instead, set the line_options metadata to 'None' (as a string).
if symbols:
plotter = this_axis.scatter
for line in range(0, num_lines):
if colors is not None:
color = colors[line]
else:
color = None
if markers is not None:
marker = markers[line]
else:
marker = None
if marker_sizes is not None:
# Note: scaling of marker sizes in scatter plots and line plots is different!
# For line plot and scatter plot marker sizes to match, the line plot
# marker size should be the square root of the scatter plot marker size.
# Maybe that should be enforced here....???
if symbols:
line_options['s'] = marker_sizes[line]
else:
line_options['markersize'] = marker_sizes[line]
if symbols:
this_line_style='None'
else:
this_line_style=line_style[line]
if marker_every is not None:
line_options['markevery'] = marker_every[line]
this_line = plotter(var_times, var_data.y[time_idxs] if num_lines == 1 else var_data.y[time_idxs, line], color=color,
linestyle=this_line_style, linewidth=thick[line], marker=marker, **line_options)
if labels is not None:
try:
if isinstance(this_line, list):
this_line[0].set_label(labels[line])
else:
this_line.set_label(labels[line])
except IndexError:
continue
if labels is not None:
legend = this_axis.legend(loc=legend_location, fontsize=legend_size, shadow=legend_shadow, title=legend_title,
labelcolor=legend_color, markerfirst=legend_markerfirst, markerscale=legend_markerscale,
facecolor=legend_facecolor, edgecolor=legend_edgecolor, frameon=legend_frameon, ncols=legend_ncols,
title_fontsize=legend_titlesize, bbox_to_anchor=bbox_to_anchor)
try:
handles = legend.legend_handles
except AttributeError:
handles = legend.legendHandles
for legobj in handles:
legobj.set_linewidth(legend_linewidth)
return True
def get_trace_options(parent_array, start_trace=None, num_traces=1, repeat=False, fill=False, fillval=None):
""" Get options for a set of traces from a parent array, extending or slicing as necessary to handle pseudovariable options
Parameters
-----------
parent_array: str or list of str
An array of option values to select from
start_trace: int
If set, we are processing a pseuodovariable, and this is the count of line traces processed so far for
previous sub-variables in the current pseuodovariable.
If parent_array is long enough (e.g. if set on the pseudovariable and intended to cover the complete
set of traces), we take a slice from start_trace with num_traces entries. If it matches
the number of traces requested (e.g. parent array comes from this subvariable), we take it as-is.
If there are fewer entries than num_traces (e.g. a single value intended to apply to all traces),
we repeat parent_array or add fill until there are enough values to take a slice from start_trace.
If None, we're just processing a regular variable, so we take values starting at zero (extending or filling
as necessary)
Default: None
num_traces: int
The number of traces in the current (sub-)variable.
repeat: bool
If True, extend the parent array by repetition if necessary. Defaults to False.
fill: bool
If True, extend the parent array by adding fill values. Defaults to False.
fillval: Any
If fill=True, values to append to parent_array to make enough entries. Defaults to None.
Returns:
--------
list of option values with num_traces entries
"""
if not isinstance(parent_array,list):
parent_array=[parent_array]
parent_length = len(parent_array)
output_array = parent_array
if start_trace is not None:
end_trace = start_trace + num_traces
if parent_length >= start_trace+num_traces:
output_array = parent_array[start_trace:end_trace]
elif parent_length == num_traces:
output_array = parent_array
else:
if repeat:
expansion_factor = int((end_trace/parent_length + 1))
expanded_array = np.tile(parent_array,expansion_factor)
output_array = expanded_array[start_trace:end_trace]
elif fill:
output_array = parent_array
missing=num_traces-parent_length
output_array.extend(np.tile([fillval],missing))
else:
logging.warning("Length of trace options (%d) smaller than number of traces (%d)",parent_length, num_traces)
else:
if len(parent_array) >= num_traces:
output_array = parent_array[0:num_traces]
else:
if repeat:
expansion_factor = int(num_traces/parent_length + 1)
expanded_array = np.tile(parent_array,expansion_factor)
output_array = expanded_array[0:num_traces]
elif fill:
output_array = parent_array
missing = num_traces - parent_length
output_array.extend(np.tile([fillval], missing))
else:
logging.warning("Length of trace options (%d) smaller than number of traces (%d)",parent_length, num_traces)
return output_array