"""Convenience module for making matplotlib plots
Your best friend in this module is the :meth:`plot` function, but the
best reference for all quirky options is :meth:`plot2d_field`.
Note:
You can't set rc parameters for this module!
"""
# FIXME: this module is way too long
from __future__ import print_function
from datetime import datetime
from distutils.version import LooseVersion
from itertools import count
import matplotlib
from matplotlib import rcParams
# hack for graceful fallback of Qt[45]Agg backends
__backend = rcParams.get('backend', None)
if __backend.lower().startswith(('qt4', 'qt5')):
try:
import matplotlib.pyplot as plt
except ImportError:
if __backend.lower().startswith('qt4'):
__new_backend = __backend.lower().replace('qt4', 'qt5')
else:
__new_backend = __backend.lower().replace('qt5', 'qt4')
try:
matplotlib.use(__new_backend, force=True, warn=False)
import matplotlib.pyplot as plt
except ImportError:
matplotlib.use(__backend, force=True, warn=False)
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from matplotlib.colors import Normalize, LogNorm, ListedColormap
import matplotlib.dates as mdates
try:
from mpl_toolkits.basemap import Basemap # pylint: disable=no-name-in-module
_HAS_BASEMAP = True
except ImportError:
_HAS_BASEMAP = False
import viscid
from viscid import coordinate
from viscid.compat import izip, string_types
from viscid import logger
from viscid import pyeval
from viscid import vutil
from viscid.plot import mpl_style # pylint: disable=unused-import
from viscid.plot import mpl_extra
from viscid.plot.mpl_direct_label import apply_labels
from viscid.plot.from_seaborn import despine
from viscid.plot import vseaborn
__mpl_ver__ = matplotlib.__version__
vseaborn.activate_from_viscid()
[docs]def plot(fld, selection=Ellipsis, force_cartesian=False, **kwargs):
"""Plot a field by dispatching to the most appropiate funciton
* If fld has 1 spatial dimensions, call
:meth:`plot1d_field(fld[selection], **kwargs)`
* If fld has 2 spatial dimensions, call
:meth:`plot2d_field(fld[selection], **kwargs)`
* If fld is 2-D and has spherical coordinates (as is the case for
ionospheric fields), try to use :meth:`plot2d_mapfield` which
uses basemap to make its axes.
Parameters:
fld (Field): Some Field
selection (optional): something that describes a field slice
force_cartesian (bool): if false, then spherical plots will use
plot_mapfield
**kwargs: Passed on to plotting function
Returns:
tuple: (plot, colorbar)
plot: matplotlib plot object
colorbar: matplotlib colorbar object
See Also:
* :meth:`plot1d_field`: target for 1d fields
* :meth:`plot2d_mapfield`: target for 2d spherical fields
* :meth:`plot2d_field`: target for 2d fields
Note:
Field slices are done using "slice_reduce", meaning extra
dimensions are reduced out.
Raises:
TypeError: Description
ValueError: Description
"""
fld = fld.slice_reduce(selection)
if not hasattr(fld, "patches"):
raise TypeError("Selection '{0}' sliced away too many "
"dimensions".format(selection))
if fld.nr_comps > 1:
raise TypeError("Scalar Fields only")
patch0 = fld.patches[0]
nr_sdims = patch0.nr_sdims
if nr_sdims == 1:
return plot1d_field(fld, **kwargs)
elif nr_sdims == 2:
is_spherical = patch0.is_spherical()
if is_spherical and not force_cartesian:
return plot2d_mapfield(fld, **kwargs)
else:
return plot2d_field(fld, **kwargs)
else:
raise ValueError("mpl can only do 1-D or 2-D fields. Either slice the "
"field yourself, or use the selection keyword "
"argument")
[docs]def plot_opts_to_kwargs(plot_opts, plot_kwargs):
"""Turn plot options from string to items in plot_kwargs
The Reason for this to to be able to specify arbitrary plotting
kwargs from the command line
Args:
plot_opts (str): plot kwargs as string
plot_kwargs (dict): kwargs to be popped by hand or passed to
plotting function
Returns:
None, the plot_opts are stuffed into plot_kwargs
"""
plot_kwargs = dict(plot_kwargs)
if not plot_opts:
return plot_kwargs
plot_opts = plot_opts.strip()
if plot_opts[0] == '{' and plot_opts[-1] == '}':
try:
import yaml
d = yaml.safe_load(plot_opts)
# if an option is given without a value, Yaml defaults to
# None, but it was probably a flag, so turn None -> True
for k in list(d.keys()):
if d[k] is None:
d[k] = True
plot_kwargs.update(d)
except ImportError:
raise ImportError("You gave plot options in YAML syntax, but "
"PyYaml is not installed. Either install PyYaml "
"or use the old comma/underscore syntax.")
else:
plot_opts = plot_opts.split(",")
for opt in plot_opts:
opt = opt.replace("=", "_").split("_")
opt[0] = opt[0].strip()
opt[1:] = [pyeval.parse(o) for o in opt[1:]]
if len(opt) == 0 or opt == ['']:
continue
elif len(opt) == 1:
plot_kwargs[opt[0]] = True
elif len(opt) == 2:
plot_kwargs[opt[0]] = opt[1]
else:
# if opt[1:] are all strings, re-combine them since some legit
# options have underscores in them, like reversed cmap names
try:
opt[1:] = "_".join(opt[1:])
except TypeError:
pass
plot_kwargs[opt[0]] = opt[1:]
return plot_kwargs
def _pop_axis_opts(plot_kwargs, default='none'):
if 'equalaxis' in plot_kwargs:
viscid.logger.warning("equalaxis option deprecated, please specify "
"this axis options explicitly (i.e., axis='equal', "
"axis='image', axis='auto', etc.)")
if 'axis' in plot_kwargs:
viscid.logger.warning("Clobbering axis option with deprecated "
"equalaxis option :'(")
if plot_kwargs.pop('equalaxis'):
plot_kwargs['axis'] = 'equal'
else:
plot_kwargs['axis'] = 'none'
using_default_viscid_axis = 'axis' not in plot_kwargs
_axis = plot_kwargs.pop("axis", default)
if _axis is not None:
_axis = _axis.strip().lower()
if _axis in ('none', ''):
_axis = None
# print('_axis =', _axis)
return _axis, using_default_viscid_axis, plot_kwargs
def _are_xy_shared(axis):
x_is_shared = len(axis.get_shared_x_axes().get_siblings(axis)) > 1
y_is_shared = len(axis.get_shared_y_axes().get_siblings(axis)) > 1
return x_is_shared, y_is_shared
def _ax_on_edge(axis):
axes_locator = axis.get_axes_locator()
if axes_locator is None:
bottom_most = axis.is_last_row()
left_most = axis.is_first_col()
else:
try:
rc_spec = axes_locator.get_subplotspec().get_rows_columns()
nrows, ncols, row_start, row_stop, col_start, col_stop = rc_spec
except AttributeError:
subplot_spec = axes_locator.get_subplotspec()
gridspec = subplot_spec.get_gridspec()
nrows, ncols = gridspec.get_geometry()
row_start, col_start = divmod(subplot_spec.num1, ncols)
if subplot_spec.num2 is not None:
row_stop, col_stop = divmod(subplot_spec.num2, ncols)
else:
row_stop = row_start
col_stop = col_start
bottom_most = row_stop >= nrows - 1
left_most = col_start == 0
return bottom_most, left_most
def _extract_actions_and_norm(axis, plot_kwargs, defaults=None):
"""
Some plot options will want to call a function after the plot is
made, like setting xlim and the like. Those are 'actions'.
Args:
axis: matplotlib axis for current plot
plot_kwargs (dict): kwargs dict containing all the options for
the current plot
defaults (dict): default values to merge plot_kwargs into
Returns:
(actions, norm_dict)
actions: list of tuples... the first erement of the tuple is
a function to call, while the 2nd is a list of arguments
to unpack when calling that function
norm_dict: will look like {'crdscale': 'lin'|'log',
'vscale': 'lin'|'log|none',
'clim': [None|number, None|number],
'symmetric': True|False}
"""
for k, v in defaults.items():
if k not in plot_kwargs:
plot_kwargs[k] = v
actions = []
if not axis:
axis = plt.gca()
if 'axis' in plot_kwargs:
_axis = plot_kwargs.pop('axis')
if _axis is not None:
if _axis == 'image':
# this is a hack to allow image axes even when we're sharing
# the x/y axes
actions.append((axis.autoscale_view, [], dict(tight=True)))
actions.append((axis.set_autoscale_on, False))
x_is_shared, y_is_shared = _are_xy_shared(axis)
if x_is_shared ^ y_is_shared:
adjustable = 'datalim'
elif (LooseVersion(__mpl_ver__) < LooseVersion("2.2")
and (x_is_shared or y_is_shared)):
adjustable = 'box-forced'
else:
adjustable = 'box'
actions.append((axis.set_aspect, 'equal',
dict(adjustable=adjustable, anchor='C')))
else:
actions.append((axis.axis, _axis))
if "x" in plot_kwargs:
actions.append((axis.set_xlim, plot_kwargs.pop('x')))
if "y" in plot_kwargs:
actions.append((axis.set_ylim, plot_kwargs.pop('y')))
if "own" in plot_kwargs:
opt = plot_kwargs.pop('own')
logger.warning("own axis doesn't seem to work yet...")
if "ownx" in plot_kwargs:
opt = plot_kwargs.pop('ownx')
logger.warning("own axis doesn't seem to work yet...")
if "owny" in plot_kwargs:
opt = plot_kwargs.pop('owny')
logger.warning("own axis doesn't seem to work yet...")
norm_dict = {'crdscale': 'lin',
'vscale': 'lin',
'clim': [None, None],
'symmetric': False
}
if plot_kwargs.pop('logscale', False):
norm_dict['vscale'] = 'log'
if "clim" in plot_kwargs:
clim = plot_kwargs.pop('clim')
norm_dict['clim'][:len(clim)] = clim
if "vmin" in plot_kwargs:
norm_dict['clim'][0] = plot_kwargs.pop('vmin')
if "vmax" in plot_kwargs:
norm_dict['clim'][1] = plot_kwargs.pop('vmax')
sym = plot_kwargs.pop('symmetric', False)
sym = plot_kwargs.pop('sym', False) or sym
norm_dict['symmetric'] = sym
# parse shorthands for specifying color scale
if "lin" in plot_kwargs:
opt = plot_kwargs.pop('lin')
norm_dict['vscale'] = 'lin'
if opt == 0:
norm_dict['symmetric'] = True
elif opt is not True:
if not isinstance(opt, (list, tuple)):
opt = [opt]
norm_dict['clim'][:len(opt)] = opt
if "log" in plot_kwargs:
opt = plot_kwargs.pop('log')
norm_dict['vscale'] = 'log'
if opt is not True:
if not isinstance(opt, (list, tuple)):
opt = [opt]
norm_dict['clim'][:len(opt)] = opt
if "loglog" in plot_kwargs:
opt = plot_kwargs.pop('loglog')
norm_dict['crdscale'] = 'log'
norm_dict['vscale'] = 'log'
if opt is not True:
if not isinstance(opt, (list, tuple)):
opt = [opt]
norm_dict['clim'][:len(opt)] = opt
# replace 'None' or 'none' with None in clim, this is kinda hacky, non?
for i in range(len(norm_dict['clim'])):
if norm_dict['clim'][i] in ["None", "none"]:
norm_dict['clim'][i] = None
# hack so that the value axis is not rescaled
if plot_kwargs.pop('norescale', False):
norm_dict['vscale'] = None
return actions, norm_dict
def _apply_actions(acts):
for act in acts:
act_args = act[1]
if not isinstance(act_args, (list, tuple)):
act_args = [act_args]
try:
act_kwargs = act[2]
except IndexError:
act_kwargs = dict()
act[0](*act_args, **act_kwargs)
def _prepare_time_axes(ax, ax_arrs, datefmt, timefmt, actions,
using_default_viscid_axis):
new_ax_arrs = [None] * len(ax_arrs)
datetime_fmt = [None] * len(ax_arrs)
for i, XI in enumerate(ax_arrs):
if viscid.is_datetime_like(XI):
# new_ax_arrs[i] = mdates.date2num(viscid.as_datetime(XI))
new_ax_arrs[i] = viscid.as_datetime(XI)
datetime_fmt[i] = datefmt
elif viscid.is_timedelta_like(XI):
# new_ax_arrs[i] = mdates.date2num(viscid.as_datetime(XI))
new_ax_arrs[i] = viscid.as_datetime(XI)
datetime_fmt[i] = timefmt
else:
new_ax_arrs[i] = XI
# with time axes, you probably don't want imageax, but only override
# this if the user didn't specify with a plot_opt
if using_default_viscid_axis and any(datetime_fmt):
for i in reversed(range(len(actions))):
if actions[i][0] in (ax.axis, ax.set_aspect):
actions.pop(i)
# take x and y plot opts and convert them to datetimes
for i, setter in enumerate([ax.set_xlim, ax.set_ylim]):
if datetime_fmt[i]:
for iact, act in enumerate(actions):
if act[0] == setter:
# print("OVERRIDE::", act)
datetimeified = viscid.as_datetime(act[1]).tolist()
actions[iact] = (setter, datetimeified)
# print(" --", actions[iact])
return new_ax_arrs, datetime_fmt
def _apply_time_axes(fig, ax, datetime_fmt, autofmt_xdate):
for fmt, axis_i in zip(datetime_fmt, (ax.xaxis, ax.yaxis)):
if fmt:
datetime_formatter = mdates.DateFormatter(fmt)
axis_i.set_major_formatter(datetime_formatter)
if axis_i is ax.xaxis and autofmt_xdate:
fig.autofmt_xdate()
def _apply_axfmt(ax, majorfmt=None, minorfmt=None, majorloc=None, minorloc=None,
which_axes="xy"):
ax_axes = {'x': ax.xaxis, 'y': ax.yaxis}
if majorfmt == "steve":
majorfmt = mpl_extra.steve_axfmt
if minorfmt == "steve":
minorfmt = mpl_extra.steve_axfmt
for axis_name in which_axes:
_axis = ax_axes[axis_name]
if majorfmt:
_axis.set_major_formatter(majorfmt)
if minorfmt:
_axis.set_minor_formatter(minorfmt)
if majorloc:
_axis.set_major_locator(majorloc)
if minorloc:
_axis.set_minor_locator(minorloc)
def _plot2d_single(ax, fld, style, namex, namey, mod, scale,
masknan, latlon, flip_plot, patchec, patchlw, patchaa,
datefmt, timefmt, autofmt_xdate, all_masked, extra_args,
actions, using_default_viscid_axis, **kwargs):
"""Make a 2d plot of a single patch
Returns:
result of the actual matplotlib plotting command
(pcolormesh, contourf, etc.)
"""
assert fld.nr_patches == 1
# pcolor mesh uses node coords, and cell data, if we have
# node data, fake it by using cell centered coords and
# trim the edges of the data... maybe i should just be
# extapolating the crds and keeping the edges...
if style in ["pcolormesh", "pcolor"]:
fld = fld.as_cell_centered()
X, Y = fld.get_crds_nc((namex, namey))
else:
if fld.iscentered("Node"):
X, Y = fld.get_crds_nc((namex, namey))
elif fld.iscentered("Cell"):
X, Y = fld.get_crds_cc((namex, namey))
# this is a hack to get rid or the white space
# between patches when contouring
Xnc, Ync = fld.get_crds_nc((namex, namey))
X[0], X[-1] = Xnc[0], Xnc[-1]
Y[0], Y[-1] = Ync[0], Ync[-1]
if latlon:
# translate latitude from 0..180 to -90..90
X, Y = np.meshgrid(X, 90 - Y)
if mod:
X *= mod[0]
Y *= mod[1]
dat = fld.data.T
if scale is not None:
dat *= scale
if masknan:
dat = np.ma.masked_where(np.isnan(dat), dat)
all_masked = all_masked and dat.mask.all()
if kwargs.pop('logscale_mask_neg', False):
dat = np.ma.masked_where(dat <= 0.0, dat)
all_masked = all_masked and dat.mask.all()
# Field.data is now xyz as are the crds
if flip_plot:
X, Y = Y.T, X.T
dat = dat.T
namex, namey = namey, namex
ax_arrs, datetime_fmt = _prepare_time_axes(ax, [X, Y], datefmt, timefmt,
actions, using_default_viscid_axis)
X, Y = ax_arrs
# datetime_fmt = [False, False]
# _XY = [X, Y]
# for i, XI in enumerate(_XY):
# if viscid.is_datetime_like(XI):
# _XY[i] = mdates.date2num(viscid.as_datetime(XI))
# datetime_fmt[i] = datefmt
# elif viscid.is_timedelta_like(XI):
# _XY[i] = mdates.date2num(viscid.as_datetime(XI))
# datetime_fmt[i] = timefmt
# X, Y = _XY
if style == "pcolormesh":
p = ax.pcolormesh(X, Y, dat, *extra_args, **kwargs)
elif style == "contour":
p = ax.contour(X, Y, dat, *extra_args, **kwargs)
elif style == "contourf":
p = ax.contourf(X, Y, dat, *extra_args, **kwargs)
elif style == "pcolor":
p = ax.pcolor(X, Y, dat, *extra_args, **kwargs)
else:
raise RuntimeError("I don't understand {0} 2d plot style".format(style))
for _fmt, axis_i in zip(datetime_fmt, (ax.xaxis, ax.yaxis)):
if _fmt:
datetime_formatter = mdates.DateFormatter(_fmt)
axis_i.set_major_formatter(datetime_formatter)
if axis_i is ax.xaxis and autofmt_xdate:
ax.get_figure().autofmt_xdate()
try:
if masknan:
p.get_cmap().set_bad(masknan)
else:
raise ValueError()
except ValueError:
p.get_cmap().set_bad('y')
# show patches?
if patchec and patchlw:
_xl = X[0]
_yl = Y[0]
_width = X[-1] - _xl
_height = Y[-1] - _yl
rect = plt.Rectangle((_xl, _yl), _width, _height,
edgecolor=patchec, linewidth=patchlw,
fill=False, antialiased=patchaa, zorder=5)
ax.add_artist(rect)
return p, all_masked
[docs]def plot2d_field(fld, ax=None, plot_opts=None, **plot_kwargs):
"""Plot a 2D Field using pcolormesh, contour, etc.
Parameters:
ax (matplotlib axis, optional): Plot in a specific axis object
plot_opts (str, optional): plot options
**plot_kwargs (str, optional): plot options
Returns:
(plot_object, colorbar_object)
See Also:
* :doc:`/plot_options`: Contains a full list of plot options
"""
patch0 = fld.patches[0]
if patch0.nr_sdims != 2:
raise RuntimeError("I will only contour a 2d field")
# raise some deprecation errors
if "extra_args" in plot_kwargs:
raise ValueError("extra_args is deprecated and for internal use only")
# init the plot by figuring out the options to use
extra_args = []
if not ax:
ax = plt.gca()
# parse plot_opts
plot_kwargs = plot_opts_to_kwargs(plot_opts, plot_kwargs)
# guess a good default value for 'axis'... i am so sorry for this logic
# but i'm kinda playing the stanly cup in a sandbox here
_xl = np.array(fld.xl[:2])
_xh = np.array(fld.xh[:2])
if 'x' in plot_kwargs:
_xl[0], _xh[0] = plot_kwargs['x']
if 'y' in plot_kwargs:
_xl[1], _xh[1] = plot_kwargs['y']
_xy_L = np.abs(_xh - _xl)
if any(viscid.is_time_like(_xli, conservative=True) for _xli in _xy_L):
_axis_def = 'none'
else:
_aspect = float(max(_xy_L)) / min(_xy_L)
_axis_def = 'image' if _aspect <= 4.0 else 'none'
# take a step back and go through any shared axes... if any of the
# shared axes use 'image', 'equal', or 'box', then this axis should too
shared_axes = set(ax.get_shared_x_axes().get_siblings(ax))
shared_axes |= set(ax.get_shared_x_axes().get_siblings(ax))
for sax in shared_axes:
_asp = sax.get_aspect()
_adj = sax.get_adjustable()
_anc = sax.get_anchor()
_aso = sax.get_autoscale_on()
if _asp == 'equal' and _adj == 'datalim' and _aso:
_axis_def = 'equal'
break
elif (_asp == 'equal' and _adj in ('box', 'box-forced') and _anc == 'C'
and not _aso):
_axis_def = 'image'
break
_axis, using_default_viscid_axis, plot_kwargs = _pop_axis_opts(plot_kwargs,
default=_axis_def)
flip_plot = plot_kwargs.pop("flipplot", False)
flip_plot = plot_kwargs.pop("flip_plot", flip_plot)
tightlim = plot_kwargs.pop('tightlim', plot_kwargs.pop('tightlims', True))
_proj = getattr(ax, "transProjection", None)
if tightlim and (_proj is None or 'identity' in _proj.__class__.__name__.lower()):
_xl = np.array(fld.xl[:2])
_xh = np.array(fld.xh[:2])
if 'x' not in plot_kwargs:
if flip_plot:
plot_kwargs['x'] = (fld.xl[1], fld.xh[1])
else:
plot_kwargs['x'] = (fld.xl[0], fld.xh[0])
if 'y' not in plot_kwargs:
if flip_plot:
plot_kwargs['y'] = (fld.xl[0], fld.xh[0])
else:
plot_kwargs['y'] = (fld.xl[1], fld.xh[1])
actions, norm_dict = _extract_actions_and_norm(ax, plot_kwargs,
defaults={'axis': _axis})
# everywhere options
scale = plot_kwargs.pop("scale", None)
masknan = plot_kwargs.pop("masknan", True)
nolabels = plot_kwargs.pop("nolabels", False)
xlabel = plot_kwargs.pop("xlabel", None)
ylabel = plot_kwargs.pop("ylabel", None)
majorfmt = plot_kwargs.pop("majorfmt", rcParams.get("viscid.majorfmt", None))
minorfmt = plot_kwargs.pop("minorfmt", rcParams.get("viscid.minorfmt", None))
majorloc = plot_kwargs.pop("majorloc", rcParams.get("viscid.majorloc", None))
minorloc = plot_kwargs.pop("minorloc", rcParams.get("viscid.minorloc", None))
datefmt = plot_kwargs.pop("datefmt", "%Y-%m-%d %H:%M:%S")
timefmt = plot_kwargs.pop("timefmt", "%H:%M:%S")
autofmt_xdate = plot_kwargs.pop("autofmtxdate", True)
autofmt_xdate = plot_kwargs.pop("autofmt_xdate", autofmt_xdate)
show = plot_kwargs.pop("show", False)
# 2d plot options
style = plot_kwargs.pop("style", "pcolormesh")
levels = plot_kwargs.pop("levels", 10)
show_grid = plot_kwargs.pop("g", False)
show_grid = plot_kwargs.pop("show_grid", show_grid)
gridec = plot_kwargs.pop("gridec", None)
gridlw = plot_kwargs.pop("gridlw", 0.25)
gridaa = plot_kwargs.pop("gridaa", True)
show_patches = plot_kwargs.pop("p", False)
show_patches = plot_kwargs.pop("show_patches", show_patches)
patchec = plot_kwargs.pop("patchec", None)
patchlw = plot_kwargs.pop("patchlw", 0.25)
patchaa = plot_kwargs.pop("patchaa", True)
mod = plot_kwargs.pop("mod", None)
title = plot_kwargs.pop("title", None)
cax = plot_kwargs.pop("cax", None)
cbar = plot_kwargs.pop("cbar", True)
colorbar = plot_kwargs.pop("colorbar", cbar)
cbar_kwargs = plot_kwargs.pop('colorbar_kwargs', dict())
cbar_kwargs = plot_kwargs.pop("cbar_kwargs", cbar_kwargs)
cbarlabel = plot_kwargs.pop("cbarlabel", None)
earth = plot_kwargs.pop("earth", False)
# undocumented options
latlon = plot_kwargs.pop("latlon", None)
norm = plot_kwargs.pop("norm", None)
action_ax = plot_kwargs.pop("action_ax", ax) # for basemap projections
# some plot_kwargs need a little more info
if show_grid:
if not isinstance(show_grid, string_types):
show_grid = 'k'
if not gridec:
gridec = show_grid
if gridec and gridlw:
plot_kwargs["edgecolors"] = gridec
plot_kwargs["linewidths"] = gridlw
plot_kwargs["antialiased"] = gridaa
if show_patches:
if not isinstance(show_patches, string_types):
show_patches = 'k'
if not patchec:
patchec = show_patches
if isinstance(colorbar, dict):
viscid.logger.warning("Deprecation, colorbar options should be passed as "
"cbar_kwargs and colorbar should be True/False.")
if cbar_kwargs:
viscid.logger.warning("Clobbering cbar_kwargs with colorbar")
colorbar, cbar_kwargs = True, colorbar
#########################
# figure out the norm...
if norm is None:
vscale = norm_dict['vscale']
vmin, vmax = norm_dict['clim']
if vmin is None:
vmin = np.nanmin([np.nanmin(blk) for blk in fld.patches])
if vmax is None:
vmax = np.nanmax([np.nanmax(blk) for blk in fld.patches])
# vmin / vmax will only be nan if all values are nan
if np.isnan(vmin) or np.isnan(vmax):
logger.warning("All-Nan encountered in Field, {0}"
"".format(patch0.name))
vmin, vmax = 1e38, 1e38
norm_dict['symmetric'] = False
if vscale == "lin":
if norm_dict['symmetric']:
maxval = max(abs(vmin), abs(vmax))
vmin = -1.0 * maxval
vmax = +1.0 * maxval
norm = Normalize(vmin, vmax)
elif vscale == "log":
if norm_dict['symmetric']:
raise ValueError("Can't use symmetric color bar with logscale")
if vmax <= 0.0:
logger.warning("Using log scale on a field with no "
"positive values")
plot_kwargs['logscale_mask_neg'] = True
vmin, vmax = 1e-20, 1e-20
elif vmin <= 0.0:
logger.warning("Using log scale on a field with values "
"<= 0. Only plotting 4 decades.")
plot_kwargs['logscale_mask_neg'] = True
vmin, vmax = vmax / 1e4, vmax
norm = LogNorm(vmin, vmax)
elif vscale is None:
norm = None
else:
raise ValueError("Unknown norm vscale: {0}".format(vscale))
if norm is not None:
plot_kwargs['norm'] = norm
else:
if isinstance(norm, Normalize):
vscale = "lin"
elif isinstance(norm, LogNorm):
vscale = "log"
else:
raise TypeError("Unrecognized norm type: {0}".format(type(norm)))
vmin, vmax = norm.vmin, norm.vmax
if "cmap" not in plot_kwargs and np.isclose(vmax, -1 * vmin, atol=0):
# by default, the symmetric_cmap is seismic (blue->white->red)
symmetric_cmap = rcParams.get("viscid.symmetric_cmap", None)
if symmetric_cmap:
plot_kwargs['cmap'] = plt.get_cmap(symmetric_cmap)
symmetric_vlims = True
else:
symmetric_vlims = False
# ok, here's some hackery for contours
if style in ["contourf", "contour"]:
if plot_kwargs.get("colors", None):
plot_kwargs['cmap'] = None
if isinstance(levels, int):
if vscale == "log":
levels = np.logspace(np.log10(vmin), np.log10(vmax), levels)
else:
levels = np.linspace(vmin, vmax, levels)
extra_args = [levels]
##############################
# now actually make the plots
namex, namey = patch0.crds.axes # fld.crds.get_culled_axes()
all_masked = False
for patch in fld.patches:
p, all_masked = _plot2d_single(action_ax, patch, style,
namex, namey, mod, scale, masknan,
latlon, flip_plot,
patchec, patchlw, patchaa, datefmt, timefmt,
autofmt_xdate, all_masked, extra_args,
actions, using_default_viscid_axis,
**plot_kwargs)
# apply option actions... this is for setting xlim / xscale / etc.
_apply_actions(actions)
if norm_dict['crdscale'] == 'log':
ax.set_xscale('log')
ax.set_yscale('log')
# figure out the colorbar...
if style == "contour" and "colors" in plot_kwargs:
colorbar = False
if masknan and all_masked:
colorbar = False
if colorbar:
if 'ax' in cbar_kwargs:
viscid.logger.warning("ignoring cbar_kwargs['ax']")
cbar_kwargs.pop('ax')
if cax is not None:
if 'cax' in cbar_kwargs:
viscid.logger.warning("clobbering colorbar['cax']")
cbar_kwargs['cax'] = cax
if "cax" in cbar_kwargs or not cbar_kwargs.get('use_grid1', True):
cax, cbar_kwargs, grid1_kwargs = _make_grid1_cbar_axes(ax, cbar_kwargs,
make_cax=False)
cax = cbar_kwargs.pop('cax', None)
else:
cax, cbar_kwargs, grid1_kwargs = _make_grid1_cbar_axes(ax, cbar_kwargs)
if "ticks" not in cbar_kwargs:
if vscale == "log":
cbar_kwargs["ticks"] = matplotlib.ticker.LogLocator()
elif symmetric_vlims:
cbar_kwargs["ticks"] = matplotlib.ticker.MaxNLocator()
else:
cbar_kwargs["ticks"] = matplotlib.ticker.LinearLocator()
cbarfmt = cbar_kwargs.pop("format", rcParams.get('viscid.cbarfmt', None))
if cbarfmt == "steve":
cbarfmt = mpl_extra.steve_cbarfmt
if cbarfmt:
cbar_kwargs["format"] = cbarfmt
cbar = plt.colorbar(p, cax=cax, **cbar_kwargs)
_cax_position = grid1_kwargs.get('position', None)
if cax and _cax_position == 'top':
cax.get_xaxis().set_ticks_position('top')
cax.get_xaxis().set_label_position('top')
# elif _cax_position == 'left':
# # cax.get_yaxis() is not a thing?
# cax.get_yaxis().set_ticks_position('left')
# cax.get_yaxis().set_label_position('left')
# apply labels... or not
if not nolabels and (cbarlabel or not title or
isinstance(title, string_types)):
if not cbarlabel:
cbarlabel = patch0.pretty_name
cbar.set_label(cbarlabel)
else:
cbar = None
if not nolabels:
# Field.data is now xyz as are the crds
if flip_plot:
namex, namey = namey, namex
if not xlabel:
xlabel = namex
if not ylabel:
ylabel = namey
if title:
if not isinstance(title, string_types):
title = patch0.pretty_name
ax.set_title(title)
x_is_shared, y_is_shared = _are_xy_shared(ax)
bottom_most, left_most = _ax_on_edge(ax)
if not x_is_shared or (x_is_shared and bottom_most):
ax.set_xlabel(namex)
if not y_is_shared or (y_is_shared and left_most):
ax.set_ylabel(namey)
_apply_axfmt(ax, majorfmt=majorfmt, minorfmt=minorfmt,
majorloc=majorloc, minorloc=minorloc)
if earth:
plot_earth(fld, axis=ax)
plt.sca(ax)
if show:
mplshow()
return p, cbar
def _make_grid1_cbar_axes(ax, cbar_kwargs, make_cax=True):
_use_grid1 = cbar_kwargs.pop('use_grid1', make_cax)
assert make_cax == _use_grid1
orig_cbar_kwargs = dict(cbar_kwargs)
position = cbar_kwargs.pop('position', None)
orientation = cbar_kwargs.pop('orientation', None)
# figure out consistant position / orientation, and warn if they are
# inconsistent
if position is None and orientation is None:
position = 'right'
orientation = 'vertical'
# sanity check orientation given position
if position in ('top', 'bottom'):
if orientation not in (None, 'horizontal'):
viscid.logger.warning("Colorbar position is '{0}', but "
"orientation '{1}' is not horizontal."
"".format(position, orientation))
orientation = 'horizontal'
if position in ('left', 'right'):
if orientation not in (None, 'vertical'):
viscid.logger.warning("Colorbar position is '{0}', but "
"orientation '{1}' is not vertical."
"".format(position, orientation))
orientation = 'vertical'
# sanity check position given orientation
if orientation == 'vertical':
if position not in ('left', 'right'):
if position is not None:
viscid.logger.warning("Colorbar orientation is horizontal, "
"but position '{0}' is neither left nor "
"right.".format(position))
position = 'right'
if orientation == 'horizontal':
if position not in ('top', 'bottom'):
if position is not None:
viscid.logger.warning("Colorbar orientation is horizontal, "
"but position '{0}' is neither left nor "
"right.".format(position))
position = 'bottom'
cbar_kwargs['orientation'] = orientation
grid1_kwargs = dict()
grid1_kwargs['orientation'] = orientation
grid1_kwargs['position'] = position
grid1_kwargs['aspect'] = cbar_kwargs.pop('aspect', 20)
default_pad = 0.05 if orientation == 'vertical' else 0.15
grid1_kwargs['pad'] = cbar_kwargs.pop('pad', default_pad)
grid1_kwargs['fraction'] = cbar_kwargs.pop('fraction', 0.05)
grid1_kwargs['shrink'] = cbar_kwargs.pop('shrink', 1.0)
cax = None
if make_cax:
try:
from viscid.plot import _mpl_grid1
cax = _mpl_grid1.make_grid1_cax(ax, **grid1_kwargs)
except ImportError:
viscid.logger.warning("Old matplotlib doesn't have "
"mpl_toolkits.axes_grid1; falling back to "
"awkward default colorbar axis.")
if cax is None:
# prepare to fallback to default mechanism, ie, let plt.colorbar
# take all the kwargs and make its own axis
cbar_kwargs = orig_cbar_kwargs
cbar_kwargs.pop('position', None)
cbar_kwargs['orientation'] = orientation
if position in ('top', 'left'):
viscid.logger.warning("Ignoring colorbar position '{0}'"
"".format(position))
ax.get_figure().sca(ax)
return cax, cbar_kwargs, grid1_kwargs
def _mlt_labels(longitude):
return "{0:g}".format(longitude * 24.0 / 360.0)
[docs]def plot2d_mapfield(fld, ax=None, plot_opts=None, **plot_kwargs):
"""Plot data on a map projection of a sphere
The default projection is polar, but any other basemap projection
can be used.
Parameters:
ax (matplotlib axis, optional): Plot in a specific axis object
plot_opts (str, optional): plot options
**plot_kwargs (str, optional): plot options
Returns:
(plot_object, colorbar_object)
Note:
Parameters are in degrees, but if the projection is 'polar',
then the plot is actually made in radians, which is important
if you want to annotate a plot.
See Also:
* :doc:`/plot_options`: Contains a full list of plot options
"""
if fld.nr_patches > 1:
raise TypeError("plot2d_mapfield doesn't do multi-patch fields yet")
# parse plot_opts
plot_kwargs = plot_opts_to_kwargs(plot_opts, plot_kwargs)
axgridec = plot_kwargs.pop("axgridec", 'grey')
axgridls = plot_kwargs.pop("axgridls", ':')
axgridlw = plot_kwargs.pop("axgridlw", 1.0)
projection = plot_kwargs.pop("projection", "polar")
hemisphere = plot_kwargs.pop("hemisphere", "none").lower().strip()
drawcoastlines = plot_kwargs.pop("drawcoastlines", False)
lon_0 = plot_kwargs.pop("lon_0", 0.0)
lat_0 = plot_kwargs.pop("lat_0", None)
bounding_lat = plot_kwargs.pop("bounding_lat", None)
bounding_lat_specified = bounding_lat is not None
title = plot_kwargs.pop("title", True)
label_lat = plot_kwargs.pop("label_lat", True)
label_mlt = plot_kwargs.pop("label_mlt", True)
# try to autodiscover hemisphere if ALL the thetas are either above
# or below the equator
lat = viscid.as_mapfield(fld, units='deg').get_crd('lat')
if hemisphere == "none":
if np.all(lat >= 0.0):
hemisphere = "north"
elif np.all(lat <= 0.0):
hemisphere = "south"
else:
viscid.logger.warning("hemisphere of field {0} is ambiguous, the "
"field contains both. Please specify either "
"north or south (defaulting to North)."
"".format(fld.name))
hemisphere = 'north'
# set a sensible default for bounding_lat... full spheres get cut off
# at 40 deg, but hemispheres or smaller are shown in full
if bounding_lat is None:
if abs(lat[-1] - lat[0]) >= 90.01:
bounding_lat = 40.0
else:
bounding_lat = 90.0
if hemisphere in ("north", 'n'):
# def_projection = "nplaea"
pass
elif hemisphere in ("south", 's'):
# def_projection = "splaea"
bounding_lat = -1 * np.abs(bounding_lat)
else:
raise ValueError("hemisphere is either 'north' or 'south'")
# boundinglat = kwargs.pop("boundinglat", def_boundinglat)
# lon_0 = kwargs.pop("lon_0", 0.0)
# lat_0 = kwargs.pop("lat_0", None)
# drawcoastlines = kwargs.pop("drawcoastlines", False)
make_periodic = plot_kwargs.get('style', None) in ("contour", "contourf")
new_fld = viscid.as_polar_mapfield(fld, bounding_lat=bounding_lat,
hemisphere=hemisphere,
make_periodic=make_periodic)
if projection != "polar" and not _HAS_BASEMAP:
viscid.logger.error("NOTE: install the basemap for the desired "
"spherical projection; falling back to "
"matplotlib's polar plot.")
projection = "polar"
if projection == "polar":
if LooseVersion(__mpl_ver__) < LooseVersion("1.1"):
raise RuntimeError("polar plots are annoying for matplotlib < ",
"version 1.1. Update your matplotlib and "
"profit.")
ax = _get_polar_axis(ax=ax)
show = plot_kwargs.pop('show', False)
plot_kwargs['nolabels'] = True
plot_kwargs['axis'] = 'none'
# hack to forceably add padding to colorbar so all labels are visible
plot_kwargs = _set_default_cbar_pad(plot_kwargs, pad=0.2)
ret = plot2d_field(new_fld, ax=ax, **plot_kwargs)
if bounding_lat_specified:
abslatlim = np.abs(bounding_lat)
else:
abslatlim = np.rad2deg(np.abs(new_fld.xh[1]))
ax.set_theta_offset(-90 * np.pi / 180.0)
if title:
if not isinstance(title, string_types):
title = new_fld.pretty_name
plt.title(title)
if axgridec:
mlt_grid_pos = (0, 45, 90, 135, 180, 225, 270, 315)
mlt_labels = (24, 3, 6, 9, 12, 15, 18, 21)
if not label_mlt:
mlt_labels = ()
ax.set_thetagrids(mlt_grid_pos, mlt_labels)
grid_label_origin = 10
if abslatlim > 50:
grid_label_dr = 20
else:
grid_label_dr = 10
# grid_dr = abs_grid_dr * np.sign(bounding_lat)
lat_grid_pos = np.arange(grid_label_origin, abslatlim, grid_label_dr)
lat_labels = np.arange(grid_label_origin, abslatlim, grid_label_dr)
if label_lat == "from_pole":
lat_labels = ["{0:g}".format(l) for l in lat_labels]
elif label_lat:
if hemisphere in ('north', 'n'):
lat_labels = 90 - lat_labels
else:
lat_labels = -90 + lat_labels
lat_labels = ["{0:g}".format(l) for l in lat_labels]
else:
lat_labels = []
ax.set_rgrids((np.pi / 180.0) * lat_grid_pos, lat_labels)
ax.set_rmax(np.deg2rad(abslatlim))
ax.grid(True, color=axgridec, linestyle=axgridls,
linewidth=axgridlw)
ax.set_axisbelow(False)
else:
ax.grid(False)
ax.set_xticklabels([])
ax.set_yticklabels([])
plt.sca(ax)
if show:
mplshow()
return ret
else:
if not ax:
ax = plt.gca()
m = Basemap(projection=projection, lon_0=lon_0, lat_0=lat_0,
boundinglat=bounding_lat, ax=ax)
show = plot_kwargs.pop('show', False)
plot_kwargs['latlon'] = True
plot_kwargs['nolabels'] = True
plot_kwargs['axis'] = 'none'
ret = plot2d_field(fld, ax=ax, action_ax=m, **plot_kwargs)
if axgridec:
if label_lat:
lat_lables = [1, 1, 1, 1]
else:
lat_lables = [0, 0, 0, 0]
if np.abs(bounding_lat) > 50.0:
latlabel_arr = np.linspace(20.0, 80.0, 4)
else:
latlabel_arr = np.linspace(50.0, 80.0, 4)
if hemisphere in ("south", 's'):
latlabel_arr *= -1
m.drawparallels(latlabel_arr, labels=lat_lables,
color=axgridec, linestyle=axgridls,
linewidth=axgridlw)
if label_mlt:
mlt_labels = [1, 1, 1, 1]
else:
mlt_labels = [0, 0, 0, 0]
m.drawmeridians(np.linspace(360.0, 0.0, 8, endpoint=False),
labels=mlt_labels, fmt=_mlt_labels,
color=axgridec, linestyle=axgridls,
linewidth=axgridlw)
if drawcoastlines:
m.drawcoastlines(linewidth=0.25)
plt.sca(ax)
if show:
mplshow()
return ret
def _set_default_cbar_pad(kwargs, pad=0.2, inplace=True):
if not inplace:
kwargs = dict(kwargs)
cbar = kwargs.get('colorbar', True)
if isinstance(cbar, dict):
set_key = 'colorbar'
cbar_kwargs = cbar
elif 'cbar_kwargs' in kwargs:
set_key = 'cbar_kwargs'
cbar_kwargs = kwargs.get('cbar_kwargs')
elif 'colorbar_kwargs' in kwargs:
set_key = 'colorbar_kwargs'
cbar_kwargs = kwargs.get('colorbar_kwargs')
else:
set_key = 'cbar_kwargs'
cbar_kwargs = dict()
if 'pad' not in cbar_kwargs:
cbar_kwargs['pad'] = pad
kwargs[set_key] = cbar_kwargs
return kwargs
[docs]def plot_iono(fld, *args, **kwargs):
"""Wrapper for easier annotated ionosphere plots
Args:
fld (Field): Some spherical field
**kwargs: Consumed, or passed to :py:func:`plot2d_mapfield`
Keyword Arugments:
scale (float): scale fld by some scalar
annotations (str): 'pot' to annotate min/max/cpcp
units (str): units for annotation values
hemisphere (str): 'north' or 'south'
fontsize (int): point of font
titlescale (float): scale fontsize of title by this much
title (str): title for the plot
Returns:
(plot_object, colorbar_object)
See Also:
* :doc:`/plot_options`: Contains a full list of plot options
Raises:
ValueError: on bad hemisphere
"""
kwargs['nolabels'] = True
scale = kwargs.pop("scale", None)
annotations = kwargs.pop("annotations", 'pot').strip().lower()
units = kwargs.pop("units", '').strip()
_fontsize = kwargs.pop("fontsize", 12)
_title_scale = kwargs.pop("titlescale", 1.25)
_title = '\n' + kwargs.get("title", fld.pretty_name).strip() + '\n'
hem = kwargs.get("hemisphere", "None").strip().lower()
fld = viscid.as_mapfield(fld, units='deg')
if scale is not None:
fld *= scale
lat = fld.get_crd('lat')
kwargs = _set_default_cbar_pad(kwargs, pad=0.2)
if units:
units = " " + units
if hem == 'none':
if np.all(lat >= 0.0):
hem = 'north'
elif np.all(lat <= 0.0):
hem = 'south'
else:
hem = 'north'
if hem in ('north', 'n'):
fldH = fld["lat=10j:"]
elif hem in ('south', 's'):
fldH = fld['theta=:-10j']
else:
raise ValueError("Unknown hemisphere: {0}".format(hem))
# now make the plot
p, cbar = plot2d_mapfield(fld, *args, **kwargs)
if annotations == 'pot':
fldH_min = np.min(fldH)
fldH_max = np.max(fldH)
plt.annotate("Min: {0:.1f}{1}".format(fldH_min, units), (-0.04, 1.0),
fontsize=_fontsize, xycoords='axes fraction')
plt.annotate("Max: {0:.1f}{1}".format(fldH_max, units), (0.75, 1.0),
fontsize=_fontsize, xycoords='axes fraction')
plt.annotate("CPCP: {0:.1f}{1}".format(fldH_max - fldH_min, units),
(-0.04, -0.04), fontsize=_fontsize, xycoords='axes fraction')
elif annotations in ('', 'none'):
pass
else:
raise ValueError("unknown annotations type: {0}".format(annotations))
plt.title(_title, fontsize=int(_title_scale * _fontsize))
# viscid.interact()
if cbar:
cbar.ax.margins(x=0.5)
return p, cbar
[docs]def plot1d_field(fld, ax=None, plot_opts=None, **plot_kwargs):
"""Plot a 1D Field using lines
Parameters:
ax (matplotlib axis, optional): Plot in a specific axis object
plot_opts (str, optional): plot options
**plot_kwargs (str, optional): plot options
See Also:
* :doc:`/plot_options`: Contains a full list of plot options
"""
patch0 = fld.patches[0]
if not ax:
ax = plt.gca()
# parse plot_opts
plot_kwargs = plot_opts_to_kwargs(plot_opts, plot_kwargs)
_axis, using_default_viscid_axis, plot_kwargs = _pop_axis_opts(plot_kwargs)
actions, norm_dict = _extract_actions_and_norm(ax, plot_kwargs,
defaults={'axis': _axis})
# everywhere options
scale = plot_kwargs.pop("scale", None)
masknan = plot_kwargs.pop("masknan", True)
nolabels = plot_kwargs.pop("nolabels", False)
xlabel = plot_kwargs.pop("xlabel", None)
ylabel = plot_kwargs.pop("ylabel", None)
majorfmt = plot_kwargs.pop("majorfmt", rcParams.get("viscid.majorfmt", None))
minorfmt = plot_kwargs.pop("minorfmt", rcParams.get("viscid.minorfmt", None))
majorloc = plot_kwargs.pop("majorloc", rcParams.get("viscid.majorloc", None))
minorloc = plot_kwargs.pop("minorloc", rcParams.get("viscid.minorloc", None))
datefmt = plot_kwargs.pop("datefmt", "%Y-%m-%d %H:%M:%S")
timefmt = plot_kwargs.pop("timefmt", "%H:%M:%S")
autofmt_xdate = plot_kwargs.pop("autofmt_xdate", True)
autofmt_xdate = plot_kwargs.pop("autofmtxdate", autofmt_xdate)
show = plot_kwargs.pop("show", False)
# 1d plot options
legend = plot_kwargs.pop("legend", False)
label = plot_kwargs.pop("label", patch0.pretty_name)
mod = plot_kwargs.pop("mod", None)
plot_kwargs["label"] = label
namex, = patch0.crds.axes
if patch0.iscentered("Node"):
x = np.concatenate([blk.get_crd_nc(namex) for blk in fld.patches])
elif patch0.iscentered("Cell"):
x = np.concatenate([blk.get_crd_cc(namex) for blk in fld.patches])
else:
raise ValueError("1d plots can do node or cell centered data only")
dat = np.concatenate([blk.data for blk in fld.patches])
ax_arrs, datetime_fmt = _prepare_time_axes(ax, [x, dat], datefmt, timefmt,
actions, using_default_viscid_axis)
x, dat = ax_arrs
if mod:
x *= mod
if scale:
dat *= scale
if masknan:
dat = np.ma.masked_where(np.isnan(dat), dat)
p = ax.plot(x, dat, **plot_kwargs)
_apply_time_axes(ax.get_figure(), ax, datetime_fmt, autofmt_xdate)
_apply_actions(actions)
###############################
# set scale based on norm_dict
vmin, vmax = norm_dict['clim']
if norm_dict['crdscale'] == 'log':
ax.set_xscale('log')
if norm_dict['vscale'] == 'log':
ax.set_yscale('log')
if norm_dict['symmetric']:
if norm_dict['vscale'] == 'log':
raise ValueError("log scale can't be symmetric about 0")
maxval = max(abs(max(dat)), abs(min(dat))) + 0.05 * (max(dat) - min(dat))
vmin, vmax = -maxval, maxval
if norm_dict['vscale'] is not None:
if vmin is not None or vmax is not None:
ax.set_ylim((vmin, vmax))
########################
# apply labels and such
if not nolabels:
if xlabel is None:
xlabel = namex
if ylabel is None:
ylabel = label
ax.set_xlabel(xlabel)
ax.set_ylabel(ylabel)
_apply_axfmt(ax, majorfmt=majorfmt, minorfmt=minorfmt,
majorloc=majorloc, minorloc=minorloc)
if legend:
if isinstance(legend, bool):
legend = 0
ax.legend(loc=0)
plt.sca(ax)
if show:
mplshow()
return p, None
[docs]def plot2d_line(line, scalars=None, **kwargs):
if scalars is not None:
scalars = [scalars]
return plot2d_lines([line], scalars=scalars, **kwargs)
[docs]def plot2d_lines(lines, scalars=None, symdir="", ax=None,
show=False, flip_plot=False, subsample=2,
pts_interp='linear', scalar_interp='linear',
marker=None, colors=None, marker_kwargs=None,
axis='none', equal=False, **kwargs):
"""Plot a list of lines in 2D
Args:
lines (list): list of 3xN ndarrays describing N xyz points
along a line
scalars (list, ndarray): a bunch of floats, rgb tuples, or
'#0000ff' colors. These can be given as one per line,
or one per vertex. See
:py:func:`viscid.vutil.prepare_lines` for more info.
symdir (str): direction perpendiclar to plane; one of 'xyz'
ax (matplotlib Axis, optional): axis on which to plot (should
be a 3d axis)
show (bool): call plt.show when finished?
flip_plot (bool): flips x and y axes on the plot
subsample (int): Factor for resampling the number of vertices.
If you are plotting a line where you want the color or line
width to change along the line to show data, you will want
`subsample > 2` to oversample the number of segments by a
factor of 2. Otherwise, the colors will be off by half a
line segment (due to limitations in matplotlib). To
undersample the lines, you can use `0 < subsample < 1`.
pts_interp (str): What kind of interpolation to use for
vertices if subsample > 0. Must be a value recognized by
:py:func:`scipy.interpolate.interp1d`.
scalar_interp (str): What kind of interpolation to use for
scalars if subsample > 0. Must be a value recognized by
:py:func:`scipy.interpolate.interp1d`.
marker (str): if given, plot the vertices using plt.scatter
colors: overrides scalar to color mapping and is passed to
matplotlib.collections.LineCollection. Note that
LineCollection only accepts rgba tuples (ie, no generic
strings). To give colors using one or more hex strings,
use `scalars='#0f0f0f'` or similar. Use `colors=zloc`
to color vertices with out-of-plane position.
marker_kwargs (dict): additional kwargs for plt.scatter
**kwargs: passed to matplotlib.collections.LineCollection
Raises:
ValueError: If a 2D plane can't be determined
Returns:
a LineCollection
"""
if not ax:
ax = plt.gca()
if isinstance(scalars, viscid.string_types) and scalars == 'zloc':
colors = 'zloc'
scalars = None
r = _prep_lines(lines, scalars=scalars, subsample=subsample,
pts_interp=pts_interp, scalar_interp=scalar_interp)
verts, segments, vert_scalars, seg_scalars, vert_colors, seg_colors, _ = r
# alpha = other['alpha']
xind, yind, zind = _xyzind_from_symdir(segments, symdir)
if flip_plot:
xind, yind = yind, xind
if colors == 'zloc':
assert zind is not None
vert_scalars = verts[zind, :]
seg_scalars = segments[:, 0, zind]
colors = None
elif colors is not None:
if seg_colors is not None:
viscid.logger.warning("plot2d_lines - overriding seg_colors with "
"explicit colors kwarg")
seg_colors = colors
line_collection = LineCollection(segments[:, :, [xind, yind]],
array=seg_scalars, colors=seg_colors,
**kwargs)
ax.add_collection(line_collection)
if marker:
if not marker_kwargs:
marker_kwargs = dict()
# if colors are not given,
if 'c' not in marker_kwargs:
if vert_colors is not None:
marker_kwargs['c'] = vert_colors
elif vert_scalars is not None:
marker_kwargs['c'] = vert_scalars
# pass along some kwargs to the scatter plot
for name in ['cmap', 'norm', 'vmin', 'vmax']:
if name in kwargs and name not in marker_kwargs:
marker_kwargs[name] = kwargs[name]
ax.scatter(verts[xind, :], verts[yind, :], marker=marker,
**marker_kwargs)
else:
_autolimit_to_vertices(ax, verts[[xind, yind], :])
if equal:
axis = 'image'
if axis.strip().lower() not in ('none', ''):
ax.axis(axis)
plt.sca(ax)
if show:
plt.show()
return line_collection
[docs]def plot3d_line(line, scalars=None, **kwargs):
if scalars is not None:
scalars = [scalars]
return plot3d_lines([line], scalars=scalars, **kwargs)
[docs]def plot3d_lines(lines, scalars=None, ax=None, show=False, subsample=2,
pts_interp='linear', scalar_interp='linear',
marker='', colors=None, marker_kwargs=None,
axis='none', equal=False, **kwargs):
"""Plot a list of lines in 3D
Args:
lines (list): list of 3xN ndarrays describing N xyz points
along a line
scalars (list, ndarray): a bunch of floats, rgb tuples, or
'#0000ff' colors. These can be given as one per line,
or one per vertex. See
:py:func:`viscid.vutil.prepare_lines` for more info.
ax (matplotlib Axis, optional): axis on which to plot (should
be a 3d axis)
show (bool): call plt.show when finished?
subsample (int): Factor for resampling the number of vertices.
If you are plotting a line where you want the color or line
width to change along the line to show data, you will want
`subsample > 2` to oversample the number of segments by a
factor of 2. Otherwise, the colors will be off by half a
line segment (due to limitations in matplotlib). To
undersample the lines, you can use `0 < subsample < 1`.
pts_interp (str): What kind of interpolation to use for
vertices if subsample > 0. Must be a value recognized by
:py:func:`scipy.interpolate.interp1d`.
scalar_interp (str): What kind of interpolation to use for
scalars if subsample > 0. Must be a value recognized by
:py:func:`scipy.interpolate.interp1d`.
marker (str): if given, plot the vertices using plt.scatter
colors: overrides scalar to color mapping and is passed to
mpl_toolkits.mplot3d.art3d.Line3DCollection. Note that this
only accepts rgba tuples (ie, no generic strings). To give
colors using one or more hex strings, use
`scalars='#0f0f0f'` or similar.
marker_kwargs (dict): additional kwargs for plt.scatter
**kwargs: passed to mpl_toolkits.mplot3d.art3d.Line3DCollection
Returns:
TYPE: Line3DCollection
"""
from mpl_toolkits.mplot3d.art3d import Line3DCollection
ax = _get_3d_axis(ax)
r = _prep_lines(lines, scalars=scalars, subsample=subsample,
pts_interp=pts_interp, scalar_interp=scalar_interp)
verts, segments, vert_scalars, seg_scalars, vert_colors, seg_colors, _ = r
if colors is not None:
if seg_colors is not None:
viscid.logger.warning("plot3d_lines - overriding seg_colors with "
"explicit colors kwarg")
seg_colors = colors
line_collection = Line3DCollection(segments[:, :, [0, 1, 2]],
array=seg_scalars, colors=seg_colors,
**kwargs)
ax.add_collection3d(line_collection)
if marker:
if not marker_kwargs:
marker_kwargs = dict()
# if colors are not given,
if 'c' not in marker_kwargs:
if vert_colors is not None:
marker_kwargs['c'] = vert_colors
elif vert_scalars is not None:
marker_kwargs['c'] = vert_scalars
# pass along some kwargs to the scatter plot
for name in ['cmap', 'norm', 'vmin', 'vmax']:
if name in kwargs and name not in marker_kwargs:
marker_kwargs[name] = kwargs[name]
ax.scatter(verts[0, :], verts[1, :], verts[2, :], marker=marker,
**marker_kwargs)
else:
_autolimit_to_vertices(ax, verts)
if equal:
axis = 'image'
if axis.strip().lower() not in ('none', ''):
ax.axis(axis)
plt.sca(ax)
if show:
plt.show()
return line_collection
[docs]def plot2d_quiver(fld, step=1, ax=None, axis='none', equal=False, **kwargs):
"""Put quivers on a 2D plot
The quivers will be plotted in the 2D plane of fld, so if fld
is 3D, then one and only one dimenstion must have shape 1.
Note:
There are some edge cases where step doesn't work.
Args:
fld(VectorField): 2.5-D Vector field to plot
step (int): only quiver every Nth grid cell. Can also
be a list of ints to specify x & y downscaling separatly
**kwargs: passed to :py:func:`matplotlpb.pyplot.quiver`
Raises:
TypeError: vector field check
ValueError: 2d field check
Returns:
result of :py:func:`matplotlpb.pyplot.quiver`
"""
if fld.nr_patches > 1:
raise TypeError("plot2d_quiver doesn't do multi-patch fields yet")
fld = fld.slice_reduce(":")
if fld.patches[0].nr_sdims != 2:
raise ValueError("2D Fields only for plot2d_quiver")
if fld.nr_comps != 3:
raise TypeError("Vector Fields only for plot2d_quiver")
# get lm axes, ie, the axes in the plane
l, m = fld.crds.axes
lm = "".join([l, m])
# get stepd scalar fields for the vector components in the plane
if not hasattr(step, "__getitem__") or len(step) < 2:
step = np.array([step, step]).reshape(-1)
first_l = (fld.shape[0] % step[0]) // 2
first_m = (fld.shape[1] % step[1]) // 2
vl = fld[l][first_l::step[0], first_m::step[1]]
vm = fld[m][first_l::step[0], first_m::step[1]]
# get coordinates
xl, xm = vl.get_crds(lm, shaped=True)
xl, xm = np.broadcast_arrays(xl, xm)
if ax is None:
ax = plt.gca()
ret = ax.quiver(xl, xm, vl, vm, **kwargs)
if equal:
axis = 'image'
if axis.strip().lower() not in ('none', ''):
ax.axis(axis)
plt.sca(ax)
return ret
[docs]def streamplot(fld, ax=None, axis='none', equal=False, **kwargs):
"""Plot 2D streamlines with :py:func:`matplotlib.pyplot.streamplot`
Args:
fld (VectorField): Some 2.5-D Vector Field
**kwargs: passed to :py:func:`matplotlib.pyplot.streamplot`
Raises:
TypeError: vector field check
ValueError: 2d field check
Returns:
result of :py:func:`matplotlib.pyplot.streamplot`
"""
if fld.nr_patches > 1:
raise TypeError("plot2d_quiver doesn't do multi-patch fields yet")
fld = fld.slice_reduce(":")
if fld.patches[0].nr_sdims != 2:
raise ValueError("2D Fields only for streamplot")
if fld.nr_comps != 3:
raise TypeError("Vector Fields only for streamplot")
# get lm axes, ie, the axes in the plane
l, m = fld.crds.axes
lm = "".join([l, m])
# get scalar fields for the vector components in the plane
fld = fld.atleast_3d()
vl, vm = fld[l], fld[m]
xl, xm = fld.get_crds(lm, shaped=False)
# matplotlib's streamplot is for uniform grids only, if crds are non
# uniform, then interpolate onto a new plane with uniform resolution
# matching the most refined region of fld
dxl = xl[1:] - xl[:-1]
dxm = xm[1:] - xm[:-1]
if not np.allclose(dxl[0], dxl) or not np.allclose(dxm[0], dxm):
# viscid.logger.warning("Matplotlib's streamplot is for uniform grids only")
vol = viscid.Volume(fld.xl, fld.xh, fld.sshape)
vl = viscid.interp_trilin(vl, vol, wrap=True)
vm = viscid.interp_trilin(vm, vol, wrap=True)
xl, xm = vl.get_crds(lm)
# interpolate linewidth and color too if given
for other in ['linewidth', 'color']:
try:
if isinstance(kwargs[other], viscid.field.Field):
o_fld = kwargs[other]
o_fld = vol.wrap_field(viscid.interp_trilin(o_fld, vol))
kwargs[other] = o_fld.slice_reduce(":")
except KeyError:
pass
# streamplot isn't happy if linewidth are color are Fields
for other in ['linewidth', 'color']:
try:
if isinstance(kwargs[other], viscid.field.Field):
kwargs[other] = kwargs[other].data.T
except KeyError:
pass
vl = vl.slice_reduce(':')
vm = vm.slice_reduce(':')
if ax is None:
ax = plt.gca()
ret = ax.streamplot(xl, xm, vl.data.T, vm.data.T, **kwargs)
if equal:
axis = 'image'
if axis.strip().lower() not in ('none', ''):
ax.axis(axis)
plt.sca(ax)
return ret
[docs]def scatter_2d(points, c='k', symdir='', flip_plot=False, ax=None, show=False,
axis='none', equal=False, **kwargs):
"""Plot scattered points on a matplotlib 3d plot
Parameters:
points: something shaped 3xN for N points, where 3 are the
xyz cartesian directions in that order
c (str, optional): color (in matplotlib format)
ax (matplotlib Axis, optional): axis on which to plot (should
be a 3d axis)
show (bool, optional): show
kwargs: passed along to :meth:`plt.statter`
"""
if not ax:
ax = plt.gca()
xind, yind, zind = _xyzind_from_symdir(points, symdir)
if flip_plot:
xind, yind = yind, xind
x = points[xind, :]
y = points[yind, :]
if c == 'zloc':
if zind is None:
raise RuntimeError("No 3rd dimension to pull colors from")
c = points[zind, :]
p = ax.scatter(x, y, c=c, **kwargs)
if equal:
axis = 'image'
if axis.strip().lower() not in ('none', ''):
ax.axis(axis)
ax.set_xlabel("xyz"[xind])
ax.set_ylabel("xyz"[yind])
plt.sca(ax)
if show:
plt.show()
return p, None
[docs]def scatter_3d(points, c='k', ax=None, show=False, axis='none', equal=False,
**kwargs):
"""Plot scattered points on a matplotlib 3d plot
Parameters:
points: something shaped 3xN for N points, where 3 are the
xyz cartesian directions in that order
c (str, optional): color (in matplotlib format)
ax (matplotlib Axis, optional): axis on which to plot (should
be a 3d axis)
show (bool, optional): show
kwargs: passed along to :meth:`plt.statter`
"""
import mpl_toolkits.mplot3d.art3d
if not ax:
ax = plt.gca(projection='3d')
x = points[0]
y = points[1]
z = points[2]
if c == 'zloc':
c = z
p = ax.scatter(x, y, z, c=c, **kwargs)
if equal:
axis = 'image'
if axis.strip().lower() not in ('none', ''):
ax.axis(axis)
ax.set_xlabel("x")
ax.set_ylabel("y")
plt.sca(ax)
if show:
plt.show()
return p, None
[docs]def tighten(**kwargs):
"""Calls `matplotlib.pyplot.tight_layout(**kwargs)`"""
try:
plt.tight_layout(**kwargs)
except AttributeError:
logger.warning("No matplotlib tight layout support")
[docs]def auto_adjust_subplots(fig=None, tight_layout=True, subplot_params=None):
"""Wrapper to adjust subplots w/ tight_layout remembering axes lims
Args:
fig (Figure): a matplotlib figure
tight_layout (bool, dict): flag for whether or not to apply a
tight layout. If a dict, then it's unpacked into
`plt.tight_layout(...)`
subplot_params (dict): unpacked into `fig.subplots_adjust(...)`
Returns:
dict: keyword arguments for fig.subplots_adjust that describe
the current figure after all adjustments are made
"""
if fig is None:
fig = plt.gcf()
# remember the axes' limits before the call to tight_layout
pre_tighten_xlim = [ax.get_xlim() for ax in fig.axes]
pre_tighten_ylim = [ax.get_ylim() for ax in fig.axes]
if tight_layout or isinstance(tight_layout, dict):
if not isinstance(tight_layout, dict):
tight_layout = {}
tighten(**tight_layout)
# apply specific subplot_params if given; hack for movies that wiggle
if subplot_params:
fig.subplots_adjust(**subplot_params)
# re-apply the old axis limits; hack for movies that wiggle
for ax, xlim, ylim in zip(fig.axes, pre_tighten_xlim, pre_tighten_ylim):
ax.set_xlim(*xlim)
ax.set_ylim(*ylim)
p = fig.subplotpars
ret = {'left': p.left, 'right': p.right, 'top': p.top,
'bottom': p.bottom, 'hspace': p.hspace, 'wspace': p.wspace}
return ret
[docs]def plot_earth(plane_spec, ax=None, scale=1.0, rot=0,
daycol='w', nightcol='k', crd_system="gse",
zorder=10, axis=None):
"""Plot a black and white Earth to show sunward direction
Parameters:
plane_spec: Specifies the plane for determining sunlit portion.
This can be a :class:`viscid.field.Field` object to try to
auto-discover the plane and crd_system, or it can be a
string like "y=0".
axis (matplotlib Axis): axis on which to plot
scale (float, optional): scale of earth
rot (float, optional): Rotation of day/night side... I forget all
the details :(
daycol (str, optional): color of dayside (matplotlib format)
nightcol (str, optional): color of nightside (matplotlib format)
crd_system (str, optional): 'mhd' or 'gse', can usually be
deduced from plane_spec if it's a Field instance.
"""
import matplotlib.patches as mpatches
# this is kind of a hacky way to
if hasattr(plane_spec, "patches"):
# this is for both Fields and AMRFields
crd_system = viscid.as_crd_system(plane_spec.patches[0], crd_system)
values = []
for blk in plane_spec.patches:
# take only the 1st reduced.nr_sdims... this should just work
try:
plane, _value = blk.meta["reduced"][0]
values.append(_value)
except KeyError:
logger.error("No reduced dims in the field, i don't know what "
"2d \nplane, we're in and can't figure out the "
"size of earth.")
return None
value = np.min(np.abs(values))
else:
plane, value = [s.strip() for s in plane_spec.split("=")]
value = float(value)
if value**2 >= scale**2:
return None
radius = np.sqrt(scale**2 - value**2)
if not ax:
if axis:
ax = axis
else:
ax = plt.gca()
if crd_system == "gse":
rot = 180
if plane == 'y' or plane == 'z':
ax.add_patch(mpatches.Wedge((0, 0), radius, 90 + rot, 270 + rot,
ec=nightcol, fc=daycol, zorder=zorder))
ax.add_patch(mpatches.Wedge((0, 0), radius, 270 + rot, 450 + rot,
ec=nightcol, fc=nightcol, zorder=zorder))
elif plane == 'x':
if value < 0:
ax.add_patch(mpatches.Circle((0, 0), radius, ec=nightcol,
fc=daycol, zorder=zorder))
else:
ax.add_patch(mpatches.Circle((0, 0), radius, ec=nightcol,
fc=nightcol, zorder=zorder))
plt.sca(ax)
return None
[docs]def get_current_colorcycle():
try:
cycle = matplotlib.rcParams['axes.prop_cycle']
return list(c['color'] for c in cycle)
except KeyError:
return list(matplotlib.rcParams['axes.color_cycle'])
[docs]def show_colorcycle(pal=None, size=1):
"""Plot the values in a color palette as a horizontal array."""
if not pal:
pal = get_current_colorcycle()
n = len(pal)
_, ax = plt.subplots(1, 1, figsize=(n * size, size))
ax.imshow(np.arange(n).reshape(1, n), cmap=ListedColormap(list(pal)),
interpolation="nearest", aspect="auto")
ax.set_xticks([])
ax.set_yticks([])
ax.set_xticklabels([])
ax.set_yticklabels([])
plt.sca(ax)
plt.show()
[docs]def show_cmap(cmap=None, size=1, aspect=8):
if cmap is None:
cmap = plt.get_cmap()
_, ax = plt.subplots(1, 1, figsize=(aspect * size, size))
gradient = np.linspace(0, 1, 256)
gradient = np.vstack((gradient, gradient))
ax.imshow(np.arange(256).reshape(1, 256), cmap=cmap,
interpolation="nearest", aspect='auto')
ax.set_xticks([])
ax.set_yticks([])
ax.set_xticklabels([])
ax.set_yticklabels([])
plt.sca(ax)
plt.show()
def _get_projected_axis(ax=None, projection='polar',
check_attr='set_thetagrids'):
_new_axis = False
if not ax:
if len(plt.gcf().axes) == 0:
_new_axis = True
ax = plt.gca()
if not hasattr(ax, check_attr):
ax = plt.subplot(*ax.get_geometry(), projection=projection,
label=str(ax.get_geometry()) + projection)
if not _new_axis:
viscid.logger.warning("Clobbering axis for subplot %s; please give a "
"%s axis if you indend to use it later.",
ax.get_geometry(), projection)
return ax
def _get_polar_axis(ax=None):
return _get_projected_axis(ax=ax, projection='polar',
check_attr='set_thetagrids')
def _get_3d_axis(ax=None):
from mpl_toolkits.mplot3d import Axes3D # pylint: disable=unused-variable
return _get_projected_axis(ax=ax, projection='3d', check_attr='zaxis')
def _autolimit_to_vertices(ax, verts):
"""Set limits on ax so that all verts are visible"""
xmin, xmax = np.min(verts[0, ...]), np.max(verts[0, ...])
ymin, ymax = np.min(verts[1, ...]), np.max(verts[1, ...])
xlim, ylim = ax.get_xlim(), ax.get_ylim()
if xlim[0] < xlim[1]:
if xmin < xlim[0]:
ax.set_xlim(left=xmin)
if xmax > xlim[1]:
ax.set_xlim(right=xmax)
else:
if xmax > xlim[0]:
ax.set_xlim(left=xmax)
if xmin < xlim[1]:
ax.set_xlim(right=xmin)
if ylim[0] < ylim[1]:
if ymin < ylim[0]:
ax.set_ylim(bottom=ymin)
if ymax > ylim[1]:
ax.set_ylim(top=ymax)
else:
if ymax > ylim[0]:
ax.set_ylim(bottom=ymax)
if ymin < ylim[1]:
ax.set_ylim(top=ymin)
# maybe z, maybe not
if verts.shape[0] > 2:
zlim = ax.get_zlim()
zmin, zmax = np.min(verts[2, ...]), np.max(verts[2, ...])
if zlim[0] < zlim[1]:
if zmin < zlim[0]:
ax.set_zlim(bottom=zmin)
if zmax > zlim[1]:
ax.set_zlim(top=zmax)
else:
if zmax > zlim[0]:
ax.set_zlim(bottom=zmax)
if zmin < zlim[1]:
ax.set_zlim(top=zmin)
def _xyzind_from_symdir(arr, symdir):
symdir = symdir.strip().lower()
if symdir:
if arr.shape[0] < 3:
raise ValueError("specifying symdir assumes arr is 3d")
if symdir == 'x':
xind, yind, zind = 1, 2, 0
elif symdir == 'y':
xind, yind, zind = 0, 2, 1
elif symdir == 'z':
xind, yind, zind = 0, 1, 2
else:
raise ValueError("symdir '{0}' not in 'xyz'".format(symdir))
else:
xy_inds = list(range(arr.shape[0]))
i = arr.shape[0] - 1
while i >= 0 and len(xy_inds) > 2:
if np.allclose(arr[i, :], arr[i, 0]):
xy_inds.pop(i)
i -= 1
if len(xy_inds) > 2:
xy_inds = [0, 1]
elif len(xy_inds) < 2:
xy_inds = [0, 0]
xind, yind = xy_inds
if arr.shape[0] > 2:
zind = min(set(range(arr.shape[0])) - set(xy_inds))
else:
zind = None
return xind, yind, zind
def _prep_lines(lines, scalars=None, subsample=2, pts_interp='linear',
scalar_interp='linear', other=None):
r = viscid.vutil.prepare_lines(lines, scalars, do_connections=True,
other=other)
verts, scalars, connections, other = r
nr_sdims = verts.shape[0]
nverts = verts.shape[1]
nsegs = connections.shape[0]
line_start = np.setdiff1d(np.arange(nverts), connections[:, 1],
assume_unique=True)
line_stop = np.concatenate([line_start[1:], [nverts]])
verts_per_line = line_stop - line_start
nlines = len(line_start)
assert nverts == nlines + nsegs
if scalars is not None:
scalars = np.atleast_2d(scalars)
else:
scalars = np.empty((0, nverts), verts.dtype)
# Use numpy / scipy to interpolate points and scalars
if subsample > 0:
fine_verts = [None] * nlines
fine_scalars = [None] * nlines
fine_connections = [None] * nlines
for i, start, stop in izip(count(), line_start, line_stop):
n_coarse = stop - start # number of verts, not segments
n_fine = int(np.ceil(subsample * (n_coarse - 1)) + 1)
coarse_verts = verts[:, start:stop]
coarse_scalars = scalars[:, start:stop]
fine_verts[i] = np.empty((nr_sdims, n_fine), dtype=verts.dtype)
fine_scalars[i] = np.empty((scalars.shape[0], n_fine),
dtype=verts.dtype)
fine_connections[i] = np.empty((n_fine - 1, 2), dtype='i')
t_coarse = np.linspace(0, 1, n_coarse)
t_fine = np.linspace(0, 1, n_fine)
try:
# raise ImportError
from scipy.interpolate import interp1d
for j in range(coarse_verts.shape[0]):
fine_verts[i][j, :] = interp1d(t_coarse, coarse_verts[j, :],
kind=pts_interp)(t_fine)
for j in range(scalars.shape[0]):
fine_scalars[i][j, :] = interp1d(t_coarse, coarse_scalars[j, :],
kind=scalar_interp)(t_fine)
except ImportError:
if pts_interp != 'linear' or scalar_interp != 'linear':
viscid.logger.error("Scipy is required to do anything "
"other than linear interpolation")
raise
for j in range(coarse_verts.shape[0]):
fine_verts[i][j, :] = np.interp(t_fine, t_coarse,
coarse_verts[j, :])
for j in range(scalars.shape[0]):
fine_scalars[i][j, :] = np.interp(t_fine, t_coarse,
coarse_scalars[j, :])
except ValueError:
# this happens in scipy's interp1d if this line has exactly 1
# vertex
fine_verts[i][j, :] = coarse_verts[j, :]
if coarse_scalars.shape[0] > 0:
fine_scalars[i][j, :] = coarse_scalars[j, :]
new_start = np.sum(np.ceil((verts_per_line[:i] - 1) * subsample) + 1)
new_stop = new_start + n_fine
fine_connections[i][:, 0] = np.arange(new_start, new_stop - 1)
fine_connections[i][:, 1] = np.arange(new_start + 1, new_stop)
verts = np.concatenate(fine_verts, axis=1)
was_uint8 = scalars.dtype == np.dtype('u1')
scalars = np.concatenate(fine_scalars, axis=1)
if was_uint8:
scalars = scalars.round().astype('u1')
connections = np.concatenate(fine_connections, axis=0)
nverts = verts.shape[1]
nsegs = connections.shape[0]
assert nsegs == nverts - nlines
# go through and make list of connected segments for the line collection
segments = np.empty((nsegs, 2, nr_sdims), dtype=verts.dtype)
segments[:, 0, :] = verts[:, connections[:, 0]].T
segments[:, 1, :] = verts[:, connections[:, 1]].T
# # TODO: straighten out array=scalars from color=scalars
colors, seg_colors = None, None
if scalars.shape[0] == 0:
scalars = None
elif scalars.shape[0] == 1:
scalars = scalars[0]
seg_scalars = None
if scalars is not None:
if scalars.dtype == np.dtype('u1'):
colors = scalars.T / 255.0
seg_colors = colors[connections[:, 0], ...]
scalars = None
else:
seg_scalars = scalars[..., connections[:, 0]]
return verts, segments, scalars, seg_scalars, colors, seg_colors, other
[docs]def interact(stack_depth=0, **kwargs):
viscid.vutil.interact(stack_depth=stack_depth + 1, mpl_ns=True, **kwargs)
# just explicitly bring in some matplotlib functions
subplot = plt.subplot
subplots = plt.subplots
[docs]def subplot2grid(*args, **kwargs):
viscid.logger.warning("pyplot.subplots should be preferred to "
"pyplot.subplot2grid")
return plt.subplot2grid(*args, **kwargs)
clf = plt.clf
savefig = plt.savefig
show = plt.show
mplshow = show
# man, i was really indecisive about these names... luckily, everything's
# a reference in Python :)
plot_line = plot3d_line
plot_lines = plot3d_lines
plot_line3d = plot3d_line
plot_lines3d = plot3d_lines
plot_line2d = plot2d_line
plot_lines2d = plot2d_lines
plot_streamlines = plot3d_lines
plot_streamlines2d = plot2d_lines
##
## EOF
##