""" Plotting routines. """
from __future__ import division
from math import log10
import numpy as np
import matplotlib.pyplot as pl
from matplotlib.collections import PolyCollection, LineCollection
import matplotlib.transforms as mtransforms
A4LANDSCAPE = 11.7, 8.3
A4PORTRAIT = 8.3, 11.7
[docs]def default_marker_size(fmt):
""" Find a default matplotlib marker size such that different marker types
look roughly the same size.
"""
temp = fmt.replace('.-', '')
if '.' in temp:
ms = 10
elif 'D' in temp:
ms = 7
elif set(temp).intersection('<>^vd'):
ms = 9
else:
ms = 8
return ms
[docs]def axvfill(xvals, ax=None, color='k', alpha=0.1, edgecolor='none', **kwargs):
""" Fill vertical regions defined by a sequence of (left, right)
positions.
Parameters
----------
xvals: list
Sequence of pairs specifying the left and right extent of each
region. e.g. (3,4) or [(0,1), (3,4)]
ax : matplotlib axes instance (default is the current axes)
The axes to plot regions on.
color : mpl colour (default 'g')
Color of the regions.
alpha : float (default 0.3)
Opacity of the regions (1=opaque).
Other keywords arguments are passed to PolyCollection.
"""
if ax is None:
ax = pl.gca()
xvals = np.asanyarray(xvals)
if xvals.ndim == 1:
xvals = xvals[None, :]
if xvals.shape[-1] != 2:
raise ValueError('Invalid input')
coords = [[(x0,0), (x0,1), (x1,1), (x1,0)] for x0,x1 in xvals]
trans = mtransforms.blended_transform_factory(ax.transData, ax.transAxes)
kwargs.update(facecolor=color, edgecolor=edgecolor, transform=trans, alpha=alpha)
p = PolyCollection(coords, **kwargs)
ax.add_collection(p)
ax.autoscale_view()
return p
[docs]def axvlines(xvals, ymin=0, ymax=1, ax=None, ls='-', color='0.7', **kwargs):
""" Plot a set of vertical lines at the given positions.
"""
if ax is None:
ax = pl.gca()
coords = [[(x,ymin), (x,ymax)] for x in xvals]
trans = mtransforms.blended_transform_factory(ax.transData, ax.transAxes)
kwargs.update(linestyle=ls, colors=color, transform=trans)
l = LineCollection(coords, **kwargs)
ax.add_collection(l)
ax.autoscale_view()
return l
[docs]def puttext(x,y,text,ax, xcoord='ax', ycoord='ax', **kwargs):
""" Print text on an axis using axes coordinates."""
if xcoord == 'data' and ycoord == 'ax':
trans = mtransforms.blended_transform_factory(ax.transData, ax.transAxes)
elif xcoord == 'ax' and ycoord == 'data':
trans = mtransforms.blended_transform_factory(ax.transAxes, ax.transData)
elif xcoord == 'ax' and ycoord == 'ax':
trans = ax.transAxes
else:
raise ValueError("Bad keyword combination: %s, %s "%(xcoord,ycoord))
return ax.text(x, y, str(text), transform=trans, **kwargs)
[docs]def distplot(vals, xvals=None, perc=(68, 95), showmean=False,
showoutliers=True, color='forestgreen', ax=None,
logx=False, logy=False, negval=None, **kwargs):
"""
Make a top-down histogram plot for an array of
distributions. Shows the median, 68%, 95% ranges and outliers.
Similar to a boxplot.
Parameters
----------
vals : sequence of arrays
2-d array or a sequence of 1-d arrays.
xvals : array of floats
x positions.
perc : array of floats (68, 95)
The percentile levels to use for area shading. Defaults show
the 68% and 95% percentile levels; roughly 1 and 2
sigma ranges for a Gaussian distribution.
showmean : boolean (False)
Whether to show the means as a dashed black line.
showoutliers : boolean (False)
Whether to show outliers past the highest percentile range.
color : mpl color ('forestgreen')
ax : mpl Axes object
Plot to this mpl Axes instance.
logx, logy : bool (False)
Whether to use a log x or y axis.
negval : float (None)
If using a log y axis, replace negative plotting values with
this value (by default it chooses a suitable value based on
the data values).
"""
if any(not hasattr(a, '__iter__') for a in vals):
raise ValueError('Input must be a 2-d array or sequence of arrays')
assert len(perc) == 2
perc = sorted(perc)
temp = 0.5*(100 - perc[0])
p1, p3 = temp, 100 - temp
temp = 0.5*(100 - perc[1])
p0, p4 = temp, 100 - temp
percentiles = p0, p1, 50, p3, p4
if ax is None:
fig = pl.figure()
ax = fig.add_subplot(111)
if xvals is None:
xvals = np.arange(len(vals), dtype=float)
# loop through columns, finding values to plot
x = []
levels = []
outliers = []
means = []
for i in range(len(vals)):
d = np.asanyarray(vals[i])
# remove nans
d = d[~np.isnan(d)]
if len(d) == 0:
# no data, skip this position
continue
# get percentile levels
levels.append(scoreatpercentile(d, percentiles))
if showmean:
means.append(d.mean())
# get outliers
if showoutliers:
outliers.append(d[(d < levels[-1][0]) | (levels[-1][4] < d)])
x.append(xvals[i])
levels = np.array(levels)
if logx and logy:
ax.loglog([],[])
elif logx:
ax.semilogx([],[])
elif logy:
ax.semilogy([],[])
if logy:
# replace negative values with a small number, negval
if negval is None:
# guess number, falling back on 1e-5
temp = levels[:,0][levels[:,0] > 0]
if len(temp) > 0:
negval = np.min(temp)
else:
negval = 1e-5
levels[~(levels > 0)] = negval
for i in range(len(outliers)):
outliers[i][outliers[i] < 0] = negval
if showmean:
if means[i] < 0:
means[i] = negval
ax.fill_between(x,levels[:,0], levels[:,1], color=color, alpha=0.2, edgecolor='none')
ax.fill_between(x,levels[:,3], levels[:,4], color=color, alpha=0.2, edgecolor='none')
ax.fill_between(x,levels[:,1], levels[:,3], color=color, alpha=0.5, edgecolor='none')
if showoutliers:
x1 = np.concatenate([[x[i]]*len(out) for i,out in enumerate(outliers)])
out1 = np.concatenate(outliers)
ax.plot(x1, out1, '.', ms=1, color='0.3')
if showmean:
ax.plot(x, means, 'k--')
ax.plot(x, levels[:,2], 'k-', **kwargs)
ax.set_xlim(xvals[0],xvals[-1])
try:
ax.minorticks_on()
except AttributeError:
pass
return ax
[docs]def errplot(x, y, yerrs, xerrs=None, fmt='.b', ax=None, ms=None, mew=0.5,
ecolor=None, elw=None, zorder=None, nonposval=None, **kwargs):
""" Plot a graph with errors.
Parameters
----------
x, y : arrays of shape (N,)
Data.
yerrs : array of shape (N,) or (N,2)
Either an array with the same length as `y`, or a list of two
such arrays, giving lower and upper limits to plot.
xerrs : array, shape (N,) or (N,2), optional
Optional x errors. The format is the same as for `yerrs`.
fmt : str
A matplotlib format string that is passed to `pylab.plot`.
ms, mew : floats
Plotting marker size and edge width.
ecolor : matplotlib color (None)
Color of the error bars. By default this will be the same color
as the markers.
elw: matplotlib line width (None)
Error bar line width.
nonposval : float (None)
Replace any non-positive values of y with `nonposval`.
"""
if ax is None:
fig = pl.figure()
ax = fig.add_subplot(111)
yerrs = np.asarray(yerrs)
if yerrs.ndim > 1:
lo = yerrs[0]
hi = yerrs[1]
else:
lo = y - yerrs
hi = y + yerrs
if nonposval is not None:
y = np.where(y <= 0, nonposval, y)
if ms is None:
ms = default_marker_size(fmt)
l, = ax.plot(x, y, fmt, ms=ms, mew=mew, **kwargs)
# find the error colour
if ecolor is None:
ecolor = l.get_mfc()
if ecolor == 'none':
ecolor = l.get_mec()
if nonposval is not None:
lo[lo <= 0] = nonposval
hi[hi <= 0] = nonposval
if 'lw' in kwargs and elw is None:
elw = kwargs['lw']
col = ax.vlines(x, lo, hi, color=ecolor, lw=elw, label='__nolabel__')
if xerrs is not None:
xerrs = np.asarray(xerrs)
if xerrs.ndim > 1:
lo = xerrs[0]
hi = xerrs[1]
else:
lo = x - xerrs
hi = x + xerrs
col2 = ax.hlines(y, lo, hi, color=ecolor, lw=elw, label='__nolabel__')
if zorder is not None:
col.set_zorder(zorder)
l.set_zorder(zorder)
if xerrs is not None:
col2.set_zorder(zorder)
if pl.isinteractive():
pl.show()
return ax
[docs]def dhist(xvals, yvals, xbins=20, ybins=20, ax=None, c='b', fmt='.', ms=1,
label=None, loc='right,bottom', xhistmax=None, yhistmax=None,
histlw=1, xtop=0.2, ytop=0.2, chist=None, **kwargs):
""" Given two set of values, plot two histograms and the
distribution.
xvals,yvals are the two properties to plot. xbins, ybins give the
number of bins or the bin edges. c is the color.
"""
if chist is None:
chist = c
if ax is None:
pl.figure()
ax = pl.gca()
loc = [l.strip().lower() for l in loc.split(',')]
if ms is None:
ms = default_marker_size(fmt)
ax.plot(xvals, yvals, fmt, color=c, ms=ms, label=label, **kwargs)
x0,x1,y0,y1 = ax.axis()
if np.__version__ < '1.5':
x,xbins = np.histogram(xvals, bins=xbins, new=True)
y,ybins = np.histogram(yvals, bins=ybins, new=True)
else:
x,xbins = np.histogram(xvals, bins=xbins)
y,ybins = np.histogram(yvals, bins=ybins)
b = np.repeat(xbins, 2)
X = np.concatenate([[0], np.repeat(x,2), [0]])
Xmax = xhistmax or X.max()
X = xtop * X / Xmax
if 'top' in loc:
X = 1 - X
trans = mtransforms.blended_transform_factory(ax.transData, ax.transAxes)
ax.plot(b, X, color=chist, transform=trans, lw=histlw)
b = np.repeat(ybins, 2)
Y = np.concatenate([[0], np.repeat(y,2), [0]])
Ymax = yhistmax or Y.max()
Y = ytop * Y / Ymax
if 'right' in loc:
Y = 1 - Y
trans = mtransforms.blended_transform_factory(ax.transAxes, ax.transData)
ax.plot(Y, b, color=chist, transform=trans, lw=histlw)
ax.set_xlim(xbins[0], xbins[-1])
ax.set_ylim(ybins[0], ybins[-1])
if pl.isinteractive():
pl.show()
return ax, dict(x=x, y=y, xbinedges=xbins, ybinedges=ybins)
[docs]def histo(a, fmt='b', bins=10, ax=None, lw=2, log=False, **kwargs):
""" Plot a histogram, without all the unnecessary stuff
matplotlib's hist() function does."""
if ax is None:
pl.figure()
ax = pl.gca()
a = np.asarray(a).ravel()
a = a[~np.isnan(a)]
vals,bins = np.histogram(a, bins=bins)
if log:
vals = np.where(vals > 0, np.log10(vals), vals)
b = np.repeat(bins, 2)
V = np.concatenate([[0], np.repeat(vals,2), [0]])
ax.plot(b, V, fmt, lw=lw, **kwargs)
if pl.isinteractive():
pl.show()
return vals,bins
[docs]def arrplot(a, x=None, y=None, ax=None, perc=(0, 100), colorbar=True,
**kwargs):
""" Plot a 2D array with coordinates.
Label coordinates such that each coloured patch representing a
value in `a` is centred on its x,y coordinate.
Parameters
----------
a : array, shape (N, M)
Values at each coordinate.
x : shape (N,)
Coordinates, must be equally spaced.
y : shape (M,)
Coordinates, must be equally spaced.
ax : axes
Axes in which to plot.
colorbar : bool (True)
Whether to also plot a colorbar.
"""
if x is None:
x = np.arange(a.shape[0])
if y is None:
y = np.arange(a.shape[1])
assert len(x) == a.shape[0]
assert len(y) == a.shape[1]
if ax is None:
pl.figure()
ax = pl.gca()
assert np.allclose(x, np.sort(x))
assert np.allclose(y, np.sort(y))
dxvals = x[1:] - x[:-1]
dx = dxvals[0]
assert np.allclose(dx, dxvals[1:])
x0, x1 = x[0] - 0.5*dx, x[-1] + 0.5*dx
dyvals = y[1:] - y[:-1]
dy = dyvals[0]
assert np.allclose(dy, dyvals[1:])
y0, y1 = y[0] - 0.5*dy, y[-1] + 0.5*dy
col = ax.imshow(a.T, aspect='auto', extent=(x0, x1, y0, y1),
interpolation='nearest', origin='lower',
**kwargs)
if colorbar:
pl.colorbar(col)
if pl.isinteractive():
pl.show()
return col
[docs]def shade_to_line(xvals, yvals, blend=1, ax=None, y0=0,
color='b'):
""" Shade a region between two curves including a color gradient.
Parameters
----------
xvals, yvals : array_like
Vertically shade to the line given by xvals, yvals
y0 : array_like
Start shading from these y values (default 0).
blend : float (default 1)
Start the cmap blending to white at this distance from `yvals`.
color : mpl color
Color used to generate the color gradient.
Returns
-------
im : mpl image object
object represeting the shaded region.
"""
if ax is None:
ax = pl.gca()
import matplotlib as mpl
yvals = np.asarray(yvals)
xvals = np.asarray(xvals)
y0 = np.atleast_1d(y0)
if len(y0) == 1:
y0 = np.ones_like(yvals) * y0[0]
else:
assert len(y0) == len(yvals)
c = [color, '1']
cm = mpl.colors.LinearSegmentedColormap.from_list('mycm', c)
ymax = yvals.max()
ymin = y0.min()
X, Y = np.meshgrid(xvals, np.linspace(ymin, ymax, 1000))
im = np.zeros_like(Y)
for i in xrange(len(xvals)):
cond = (Y[:, i] > yvals[i] - blend) & (Y[:, i] > y0[i])
im[cond, i] = (Y[cond, i] - (yvals[i] - blend)) / blend
cond = Y[:, i] > yvals[i]
im[cond, i] = 1
cond = Y[:, i] < y0[i]
im[cond, i] = 0
im = ax.imshow(im, extent=(xvals[0], xvals[-1], ymin, ymax),
origin='lower', cmap=cm, aspect='auto')
return im
[docs]def shade_to_line_vert(yvals, xvals, blend=1, ax=None, x0=0,
color='b'):
""" Shade a region between two curves including a color gradient.
Parameters
----------
yvals, xvals : array_like
horizontally shade to the line given by xvals, yvals
x0 : array_like
Start shading from these x values (default 0).
blend : float (default 1)
Start the cmap blending to white at this distance from `yvals`.
color : mpl color
Color used to generate the color gradient.
Returns
-------
im : mpl image object
object represeting the shaded region.
"""
if ax is None:
ax = pl.gca()
import matplotlib as mpl
yvals = np.asarray(yvals)
xvals = np.asarray(xvals)
x0 = np.atleast_1d(x0)
if len(x0) == 1:
x0 = np.ones_like(xvals) * x0[0]
else:
assert len(x0) == len(xvals)
c = [color, '1']
cm = mpl.colors.LinearSegmentedColormap.from_list('mycm', c)
xmax = xvals.max()
xmin = x0.min()
Y, X = np.meshgrid(yvals, np.linspace(xmin, xmax, 1000))
im = np.zeros_like(X)
for i in xrange(len(yvals)):
cond = (X[:, i] > xvals[i] - blend) & (X[:, i] > x0[i])
im[cond, i] = (X[cond, i] - (xvals[i] - blend)) / blend
cond = X[:, i] > xvals[i]
im[cond, i] = 1
cond = X[:, i] < x0[i]
im[cond, i] = 0
art = ax.imshow(im.T, extent=(xmin, xmax, yvals[0], yvals[-1]),
origin='lower', cmap=cm, aspect='auto')
return art, im, X, Y
[docs]def draw_arrows(x, y, ax=None, capsize=2, ms=6, direction='up',
c='k', **kwargs):
""" Draw arrows that can be used to show limits.
Extra keyword arguments are passed to `pyplot.scatter()`. To draw
a shorter arrow, get the arrow length desired by reducing the `ms`
value, then increase capsize until you are happy with the result,
vice versa to draw a longer arrow.
Parameters
----------
x, y: float or arrays of shape (N,)
x and y positions.
direction: str {'up', 'down', 'left', 'right'}
The direction in which the arrows should point.
"""
arrowlength=10.
capsize = min(capsize, arrowlength)
yvert = np.array([0, arrowlength, arrowlength - capsize, arrowlength,
arrowlength - capsize, arrowlength])
xvert = np.array([0, 0, 0.5*capsize, 0, -0.5*capsize, 0])
if direction == 'down':
arrow_verts = zip(xvert, -yvert)
elif direction == 'up':
arrow_verts = zip(xvert, yvert)
elif direction == 'left':
arrow_verts = zip(-yvert, xvert)
elif direction == 'up':
arrow_verts = zip(yvert, xvert)
else:
raise ValueError(
"direction must be one of 'up', 'down', 'left', 'right'")
if ax is None:
pl.figure()
ax = pl.gca()
c = ax.scatter(x, y, s=(1000/6.)*ms, marker=None, verts=arrow_verts,
edgecolors=c, **kwargs)
return c
[docs]def calc_log_minor_ticks(majticks):
""" Get minor tick positions for a log scale.
Parameters
----------
majticks : array_like
log10 of the major tick positions.
Returns
-------
minticks : ndarray
log10 of the minor tick positions.
"""
tickpos = np.log10(np.arange(2, 10))
minticks = []
for t in np.atleast_1d(majticks):
minticks.extend(t + tickpos)
return minticks
[docs]def plot_ticks_wa(ax, wa, fl, height, ticks, keeponly=None, labels=True):
""" plot a ticks on a wavelength scale.
This plots ticks (such as those returned by `find_tau()`) on a
spectrum.
Parameters
----------
ax : matplotlib axes
The axes on which to plot the ticks.
wa, fl : array_like
wavelength and flux of spectrum. `wa` must be sorted.
height : float
tick height in flux units.
ticks : record array
A record array of the sort returned by `find_tau`. The fields
wa, wa0, and name are required.
keeponly : str
If this is not None (the default), then only plot ticks that
contain this string in their name.
labels : bool (True)
Whether to plot labels next to the tickmarks.
Returns
-------
Ticks, Tlabels : Matplotlib collection of tickmarks and tick labels.
The artists corresponding to the ticks and their labels.
"""
ind = wa.searchsorted(ticks.wa)
c0 = (ind == 0) | (ind == len(wa))
ticks = ticks[~c0]
ymin = fl[ind[~c0]]*1.1
Tlabels = []
c1 = np.ones(len(ticks), bool)
for i,t in enumerate(ticks):
if keeponly is not None:
if keeponly not in t.name:
c1[i] = False
continue
if not labels:
continue
label = '%s %.0f' % (t.name, t.wa0)
label = label.replace('NeVII', 'NeVIII')
Tlabels.append(ax.text(t.wa, ymin[i] + 1.1*height, label, rotation=60,
fontsize=8, va='bottom', alpha=0.7))
Ticks = ax.vlines(ticks.wa[c1], ymin[c1], ymin[c1] + height, color='c',
lw=1)
return Ticks, Tlabels