def corner(xs,
labels=None,
extents=None,
truths=None,
truth_color="#4682b4",
scale_hist=False,
bblocks=False,
quantiles=[],
verbose=True,
plot_contours=True,
plot_datapoints=True,
fig=None,
**kwargs):
"""
Make a *sick* corner plot showing the projections of a data set in a
multi-dimensional space. kwargs are passed to hist2d() or used for
`matplotlib` styling.
Parameters
----------
xs : array_like (nsamples, ndim)
The samples. This should be a 1- or 2-dimensional array. For a 1-D
array this results in a simple histogram. For a 2-D array, the zeroth
axis is the list of samples and the next axis are the dimensions of
the space.
labels : iterable (ndim,) (optional)
A list of names for the dimensions.
extents : iterable (ndim,) (optional)
A list of length 2 tuples containing lower and upper bounds (extents)
for each dimension, e.g., [(0.,10.), (1.,5), etc.]
truths : iterable (ndim,) (optional)
A list of reference values to indicate on the plots.
truth_color : str (optional)
A ``matplotlib`` style color for the ``truths`` makers.
scale_hist : bool (optional)
Should the 1-D histograms be scaled in such a way that the zero line
is visible?
bblocks : bool (optional)
Option to bin histogram data using an adaptive-width algorithm written
by Jake VanderPlas for astroML following Scargle et al. 2012
(http://adsabs.harvard.edu/abs/2012arXiv1207.5578S) - implemented
here as in astroML's "hist" plotting function
quantiles : iterable (optional)
A list of fractional quantiles to show on the 1-D histograms as
vertical dashed lines.
verbose : bool (optional)
If true, print the values of the computed quantiles.
plot_contours : bool (optional)
Draw contours for dense regions of the plot.
plot_datapoints : bool (optional)
Draw the individual data points.
fig : matplotlib.Figure (optional)
Overplot onto the provided figure object.
"""
# Deal with 1D sample lists.
xs = np.atleast_1d(xs)
if len(xs.shape) == 1:
xs = np.atleast_2d(xs)
else:
assert len(xs.shape) == 2, "The input sample array must be 1- or 2-D."
xs = xs.T
assert xs.shape[0] <= xs.shape[1], "I don't believe that you want more " \
"dimensions than samples!"
# backwards-compatibility
plot_contours = kwargs.get("smooth", plot_contours)
K = len(xs)
factor = 2.0 # size of one side of one panel
lbdim = 0.5 * factor # size of left/bottom margin
trdim = 0.05 * factor # size of top/right margin
whspace = 0.05 # w/hspace size
plotdim = factor * K + factor * (K - 1.) * whspace
dim = lbdim + plotdim + trdim
if fig is None:
fig, axes = pl.subplots(K, K, figsize=(dim, dim))
else:
try:
axes = np.array(fig.axes).reshape((K, K))
except:
raise ValueError("Provided figure has {0} axes, but data has "
"dimensions K={1}".format(len(fig.axes), K))
lb = lbdim / dim
tr = (lbdim + plotdim) / dim
fig.subplots_adjust(left=lb,
bottom=lb,
right=tr,
top=tr,
wspace=whspace,
hspace=whspace)
if extents is None:
extents = [[x.min(), x.max()] for x in xs]
# Check for parameters that never change.
m = np.array([e[0] == e[1] for e in extents], dtype=bool)
if np.any(m):
raise ValueError(("It looks like the parameter(s) in column(s) "
"{0} have no dynamic range. Please provide an "
"`extent` argument.").format(", ".join(
map("{0}".format,
np.arange(len(m))[m]))))
if bblocks:
if not has_bb:
raise ImportError("You can't bin in bayesian blocks without " \
"installing astroML")
# have kwarg to send to hist2d
kwargs["bblocks"] = True
for i, x in enumerate(xs):
ax = axes[i, i]
# Get bin edges from bayesian_blocks
if bblocks:
x = x[(x >= extents[i][0]) & (x <= extents[i][1])]
kwargs["bins"] = bb(x)
# Plot the histograms.
n, b, p = ax.hist(x,
bins=kwargs.get("bins", 50),
range=extents[i],
histtype="step",
color=kwargs.get("color", "k"))
if truths is not None:
ax.axvline(truths[i], color=truth_color)
# Plot quantiles if wanted.
if len(quantiles) > 0:
xsorted = sorted(x)
qvalues = [xsorted[int(q * len(xsorted))] for q in quantiles]
for q in qvalues:
ax.axvline(q, ls="dashed", color=kwargs.get("color", "k"))
if verbose:
print("Quantiles:")
print(zip(quantiles, qvalues))
# Set up the axes.
ax.set_xlim(extents[i])
if scale_hist:
maxn = np.max(n)
ax.set_ylim(-0.1 * maxn, 1.1 * maxn)
else:
ax.set_ylim(0, 1.1 * np.max(n))
ax.set_yticklabels([])
ax.xaxis.set_major_locator(MaxNLocator(5))
# Not so DRY.
if i < K - 1:
ax.set_xticklabels([])
else:
[l.set_rotation(45) for l in ax.get_xticklabels()]
if labels is not None:
ax.set_xlabel(labels[i])
ax.xaxis.set_label_coords(0.5, -0.3)
for j, y in enumerate(xs):
ax = axes[i, j]
if j > i:
ax.set_visible(False)
ax.set_frame_on(False)
continue
elif j == i:
continue
hist2d(y,
x,
ax=ax,
extent=[extents[j], extents[i]],
plot_contours=plot_contours,
plot_datapoints=plot_datapoints,
**kwargs)
if truths is not None:
ax.plot(truths[j], truths[i], "s", color=truth_color)
ax.axvline(truths[j], color=truth_color)
ax.axhline(truths[i], color=truth_color)
ax.xaxis.set_major_locator(MaxNLocator(5))
ax.yaxis.set_major_locator(MaxNLocator(5))
if i < K - 1:
ax.set_xticklabels([])
else:
[l.set_rotation(45) for l in ax.get_xticklabels()]
if labels is not None:
ax.set_xlabel(labels[j])
ax.xaxis.set_label_coords(0.5, -0.3)
if j > 0:
ax.set_yticklabels([])
else:
[l.set_rotation(45) for l in ax.get_yticklabels()]
if labels is not None:
ax.set_ylabel(labels[i])
ax.yaxis.set_label_coords(-0.3, 0.5)
return fig
def hist2d(x, y, *args, **kwargs):
"""
Plot a 2-D histogram of samples.
"""
ax = kwargs.pop("ax", pl.gca())
extent = kwargs.pop("extent", [[x.min(), x.max()], [y.min(), y.max()]])
bins = kwargs.pop("bins", 50)
color = kwargs.pop("color", "k")
linewidths = kwargs.pop("linewidths", None)
plot_datapoints = kwargs.get("plot_datapoints", True)
plot_contours = kwargs.get("plot_contours", True)
cmap = cm.get_cmap("gray")
cmap._init()
cmap._lut[:-3, :-1] = 0.
cmap._lut[:-3, -1] = np.linspace(1, 0, cmap.N)
if kwargs.pop("bblocks", False):
# change binning if bayesian_blocks used
x = x[(x >= extent[0][0]) & (x <= extent[0][1])]
X = bb(x)
Y = bins # these were already calculated above
else:
X = np.linspace(extent[0][0], extent[0][1], bins + 1)
Y = np.linspace(extent[1][0], extent[1][1], bins + 1)
try:
H, X, Y = np.histogram2d(x.flatten(), y.flatten(), bins=(X, Y))
except ValueError:
raise ValueError("It looks like at least one of your sample columns "
"have no dynamic range. You could try using the "
"`extent` argument.")
V = 1.0 - np.exp(-0.5 * np.arange(0.5, 2.1, 0.5)**2)
Hflat = H.flatten()
inds = np.argsort(Hflat)[::-1]
Hflat = Hflat[inds]
sm = np.cumsum(Hflat)
sm /= sm[-1]
for i, v0 in enumerate(V):
try:
V[i] = Hflat[sm <= v0][-1]
except:
V[i] = Hflat[0]
X1, Y1 = 0.5 * (X[1:] + X[:-1]), 0.5 * (Y[1:] + Y[:-1])
X, Y = X[:-1], Y[:-1]
if plot_datapoints:
ax.plot(x,
y,
"o",
color=color,
ms=1.5,
zorder=-1,
alpha=0.1,
rasterized=True)
if plot_contours:
ax.contourf(X1,
Y1,
H.T, [V[-1], H.max()],
cmap=LinearSegmentedColormap.from_list(
"cmap", ([1] * 3, [1] * 3), N=2),
antialiased=False)
if plot_contours:
ax.pcolor(X, Y, H.max() - H.T, cmap=cmap)
ax.contour(X1, Y1, H.T, V, colors=color, linewidths=linewidths)
data = np.vstack([x, y])
mu = np.mean(data, axis=1)
cov = np.cov(data)
if kwargs.pop("plot_ellipse", False):
error_ellipse(mu, cov, ax=ax, edgecolor="r", ls="dashed")
ax.set_xlim(extent[0])
ax.set_ylim(extent[1])