import numpy as np
import warnings
from pathlib import Path
from scipy.stats import beta
from corner import corner
import matplotlib.pyplot as plt
from matplotlib import axes, colormaps
from matplotlib.projections import projection_registry
from matplotlib.lines import Line2D
from matplotlib.colors import Normalize
from matplotlib.gridspec import GridSpec
from figaro import plot_settings
from figaro.marginal import marginalise
from figaro.credible_regions import ConfidenceArea
from figaro.cumulative import fast_cumulative
from figaro.utils import recursive_grid
from figaro.mixture import mixture
# Telling python to ignore empty legend warning from matplotlib
warnings.filterwarnings("ignore", message = "No artists with labels found to put in legend. Note that artists whose label start with an underscore are ignored when legend() is called with no argument.")
[docs]
class PPPlot(axes.Axes):
"""
Construct a probability-probability (P-P) plot.
Derived from https://lscsoft.docs.ligo.org/ligo.skymap/_modules/ligo/skymap/plot/pp.html#PPPlot
This class avoids installing the whole ligo.skymap.plot package.
"""
name = 'pp_plot'
def __init__(self, *args, **kwargs):
# Call parent constructor
super().__init__(*args, **kwargs)
# Square axes, limits from 0 to 1
self.set_aspect(1.0)
self.set_xlim(0.0, 1.0)
self.set_ylim(0.0, 1.0)
[docs]
def add_diagonal(self, *args, **kwargs):
"""
Add a diagonal line to the plot, running from (0, 0) to (1, 1).
Other parameters
----------------
kwargs :
optional extra arguments to `matplotlib.axes.Axes.plot`
"""
# Make copy of kwargs to pass to plot()
kwargs = dict(kwargs)
kwargs.setdefault('color', 'black')
kwargs.setdefault('linestyle', 'dashed')
kwargs.setdefault('linewidth', 0.5)
# Plot diagonal line
return self.plot([0, 1], [0, 1], *args, **kwargs)
[docs]
def add_confidence_band(self, nsamples, cl=0.9, **kwargs):
"""
Add a target confidence band.
Parameters
----------
nsamples : int
Number of P-values
cl : float, default: 0.9
Confidence level
Other parameters
----------------
**kwargs :
optional extra arguments to `matplotlib.axes.Axes.fill_betweenx`
"""
n = nsamples
k = np.arange(0, n + 1)
p = k / n
ci_lo, ci_hi = beta.interval(cl, k + 1, n - k + 1)
# Make copy of kwargs to pass to fill_betweenx()
kwargs = dict(kwargs)
kwargs.setdefault('color', 'ghostwhite')
kwargs.setdefault('edgecolor', 'gray')
kwargs.setdefault('linewidth', 0.5)
kwargs.setdefault('alpha', 0.5)
fontsize = kwargs.pop('fontsize', 'x-small')
return self.fill_betweenx(p, ci_lo, ci_hi, **kwargs, label = '${0}\%\ CR$'.format(int(cl*100)))
@classmethod
def _as_mpl_axes(cls):
"""
Support placement in figure using the `projection` keyword argument.
See http://matplotlib.org/devel/add_new_projection.html.
"""
return cls, {}
projection_registry.register(PPPlot)
[docs]
def plot_multidim(draws, samples = None, bounds = None, out_folder = '.', name = 'density', labels = None, units = None, hierarchical = False, show = False, save = True, subfolder = False, n_pts = 200, true_value = None, levels = [0.5, 0.68, 0.9], scatter_points = False, median_label = None, fig = None, colors = ['steelblue', 'darkturquoise', 'mediumturquoise'], colormap = 'Blues'):
"""
Plot the recovered multidimensional distribution along with samples from the true distribution (if available) as corner plot.
Arguments:
iterable draws: container for mixture instances
np.ndarray samples: samples from the true distribution (if available)
iterable bounds: bounds for the recovered distribution. If None, bounds from mixture instances are used.
str or Path out_folder: output folder
str name: name to be given to outputs
list-of-str labels: LaTeX-style quantity label, for plotting purposes
list-of-str units: LaTeX-style quantity unit, for plotting purposes
bool hierarchical: hierarchical inference, for plotting purposes
bool save: whether to save the plot or not
bool show: whether to show the plot during the run or not
bool subfolder: whether to save in a dedicated subfolder
int n_pts: number of grid points (same for each dimension)
iterable true_value: true value to plot
iterable levels: credible levels to plot
bool scatter_points: scatter samples on 2d plots
str median_label: label to assign to the reconstruction
matplotlib.figure.Figure fig: figure to use for plotting. Must have (dim,dim) axes.
list-of-str colors: list of colors for median, 68% and 90% credible regions
str colormap: colormap for contour plots
Returns:
matplotlib.figure.Figure: figure with the plot
"""
dim = draws[0].dim
if median_label is None:
if hierarchical:
median_label = '\mathrm{(H)DPGMM}'
else:
median_label = '\mathrm{DPGMM}'
color_med, color_68, color_90 = colors
if labels is None:
labels = ['$x_{'+'{0}'.format(i+1)+'}$' for i in range(dim)]
else:
labels = ['${0}$'.format(l) for l in labels]
if units is not None:
labels = [l[:-1]+'\ [{0}]$'.format(u) if not u == '' else l for l, u in zip(labels, units)]
levels = np.atleast_1d(levels)
ext_bounds = False
if bounds is not None:
ext_bounds = True
all_bounds = np.atleast_2d([d.bounds for d in draws])
x_min = np.min(all_bounds, axis = -1).max(axis = 0)
x_max = np.max(all_bounds, axis = -1).min(axis = 0)
probit = np.array([d.probit for d in draws]).any()
if bounds is not None:
bounds = np.atleast_2d(bounds)
if bounds.shape == (1, 2):
bounds = np.array([bounds[0] for _ in range(dim)])
elif bounds.shape == (dim, 2):
if not (bounds[:,0] >= x_min).all() and probit:
warnings.warn("The provided lower bound is invalid for at least one draw. Default values will be used instead.")
x_min[np.where(bounds[:,0] >= x_min)] = bounds[:,0][np.where(bounds[:,0] >= x_min)]
if not (bounds[:,1] <= x_max).all() and probit:
warnings.warn("The provided upper bound is invalid for at least one draw. Default values will be used instead.")
x_max[np.where(bounds[:,1] <= x_max)] = bounds[:,1][np.where(bounds[:,1] <= x_max)]
else:
warnings.warn("The provided bounds have an invalid shape {0}. Shape must be (1,2) or (dim, 2).\nDefault bounds will be used instead.".format(bounds.shape))
bounds = np.array([x_min, x_max]).T
K = dim
factor = 3.0 # size of one side of one panel
lbdim = 0.5 * factor # size of left/bottom margin
trdim = 0.2 * factor # size of top/right margin
whspace = 0.2 # w/hspace size
plotdim = factor * dim + factor * (K - 1.0) * whspace
dim_plt = lbdim + plotdim + trdim
if fig is None:
fig, axs = plt.subplots(K, K, figsize=(dim_plt, dim_plt))
else:
axs = np.array(fig.axes).reshape(dim,dim)
# Format the figure.
lb = lbdim / dim_plt
tr = (lbdim + plotdim) / dim_plt
fig.subplots_adjust(left=lb, bottom=lb, right=tr, top=tr, wspace=whspace, hspace=whspace)
# Samples (if available)
if samples is not None:
bins = np.array([int(np.sqrt(len(samples[:, c]))) for c in range(dim)])
corner(samples, color = '#1f77b4', fig = fig, hist_kwargs = {'density': True, 'label':'$\mathrm{Samples}$', 'linewidth':0.7} , plot_density = False, contour_kwargs = {'linewidths':0.3, 'linestyles':'dashed'}, levels = [0.5,0.68,0.9], no_fill_contours = True, hist_bin_factor = bins/20, quiet = True)
# 1D plots (diagonal)
for column in range(K):
ax = axs[column, column]
# Marginalise over all uninterested columns
dims = list(np.arange(dim))
dims.remove(column)
try:
if np.all([isinstance(d, mixture) for d in draws]):
marg_draws = marginalise(draws, dims)
else:
marg_draws = np.array([d.marginalise(dims) for d in draws])
except AttributeError:
raise AttributeError("The provided instances do not have a marginalise() method")
# Credible regions
lim = bounds[column]
if samples is not None and not ext_bounds:
lim_l = ax.get_xlim()
lim[0] = np.max((lim[0], lim_l[0]))
lim[1] = np.min((lim[1], lim_l[1]))
x = np.linspace(lim[0], lim[1], n_pts)
dx = x[1]-x[0]
probs = np.array([d.pdf(x) for d in marg_draws])
percentiles = [50, 5, 16, 84, 95]
p = {}
for perc in percentiles:
p[perc] = np.percentile(probs, perc, axis = 0)
norm = p[50].sum()*dx
for perc in percentiles:
p[perc] = p[perc]/norm
# CR
ax.fill_between(x, p[95], p[5], color = color_90, alpha = 0.25)
ax.fill_between(x, p[84], p[16], color = color_68, alpha = 0.25)
if true_value is not None:
if true_value[column] is not None:
ax.axvline(true_value[column], c = 'firebrick', lw = 0.5)
ax.plot(x, p[50], lw = 0.7, color = color_med, label = '${0}$'.format(median_label))
ax.set_yticks([])
if column == K - 1:
if labels is not None:
ax.set_xlabel(labels[-1])
ticks = np.linspace(lim[0], lim[1], 5)
ax.set_xticks(ticks)
[l.set_rotation(45) for l in ax.get_xticklabels()]
ax.set_xlim(lim[0], lim[1])
if column < K-1:
ax.set_xticklabels([])
ax.set_yticklabels([])
# 2D plots (off-diagonal)
for row in range(K):
for column in range(K):
ax = axs[row,column]
if column > row:
ax.set_frame_on(False)
ax.set_xticks([])
ax.set_yticks([])
continue
elif column == row:
continue
# Marginalise
dims = list(np.arange(dim))
dims.remove(column)
dims.remove(row)
try:
if np.all([isinstance(d, mixture) for d in draws]):
marg_draws = marginalise(draws, dims)
else:
marg_draws = np.array([d.marginalise(dims) for d in draws])
except:
raise AttributeError("The provided instances do not have a marginalise() method")
# Credible regions
lim = bounds[[row, column]]
if samples is not None and not ext_bounds:
lim_l = np.array([ax.get_ylim(), ax.get_xlim()])
for i in range(2):
lim[i,0] = np.max((lim[i,0], lim_l[i,0]))
lim[i,1] = np.min((lim[i,1], lim_l[i,1]))
grid, dgrid = recursive_grid(lim[::-1], np.ones(2, dtype = int)*int(n_pts))
x = np.linspace(lim[0,0], lim[0,1], n_pts)
y = np.linspace(lim[1,0], lim[1,1], n_pts)
dd = np.array([d.pdf(grid) for d in marg_draws])
median = np.percentile(dd, 50, axis = 0)
median = median/(median.sum()*np.prod(dgrid))
median = median.reshape(n_pts, n_pts)
X,Y = np.meshgrid(x,y)
with np.errstate(divide = 'ignore'):
logmedian = np.nan_to_num(np.log(median), nan = -np.inf, neginf = -np.inf)
_,_,levs = ConfidenceArea(logmedian, x, y, adLevels=levels)
ax.contourf(Y, X, np.exp(logmedian), cmap = colormap, levels = 100)
if true_value is not None:
if true_value[row] is not None:
ax.axhline(true_value[row], c = 'firebrick', lw = 0.5)
if true_value[column] is not None:
ax.axvline(true_value[column], c = 'firebrick', lw = 0.5)
if true_value[column] is not None and true_value[row] is not None:
ax.plot(true_value[column], true_value[row], color = 'firebrick', marker = 's', ms = 3)
c1 = ax.contour(Y, X, logmedian, np.sort(levs), colors=color_med, linewidths=0.3)
if plot_settings.tex_flag:
ax.clabel(c1, fmt = {l:'{0:.0f}\\%'.format(100*s) for l,s in zip(c1.levels, np.sort(levels)[::-1])}, fontsize = 3)
else:
ax.clabel(c1, fmt = {l:'{0:.0f}\%'.format(100*s) for l,s in zip(c1.levels, np.sort(levels)[::-1])}, fontsize = 3)
# Samples (if available)
if samples is not None and scatter_points:
ax.scatter(samples[:,column], samples[:,row], marker = '+', c = 'firebrick', linewidths = 1)
# Ticks
xticks = np.linspace(lim[1,0], lim[1,1], 5)
ax.set_xticks(xticks)
yticks = np.linspace(lim[0,0], lim[0,1], 5)
ax.set_yticks(yticks)
ax.set_xlim(*lim[1])
ax.set_ylim(*lim[0])
if column == 0:
ax.set_ylabel(labels[row])
[l.set_rotation(45) for l in ax.get_yticklabels()]
else:
ax.set_yticklabels([])
if row == K - 1:
[l.set_rotation(45) for l in ax.get_xticklabels()]
ax.set_xlabel(labels[column])
else:
ax.set_xticklabels([])
fig.axes[K-1].legend(*fig.axes[0].get_legend_handles_labels(), loc = 'center')
fig.align_labels()
if show:
plt.show()
if save:
if not subfolder:
fig.savefig(Path(out_folder, '{0}.pdf'.format(name)), bbox_inches = 'tight')
else:
if not Path(out_folder, 'density').exists():
try:
Path(out_folder, 'density').mkdir()
except FileExistsError:
pass
fig.savefig(Path(out_folder, 'density', '{0}.pdf'.format(name)), bbox_inches = 'tight')
if (save or show):
plt.close()
return fig
[docs]
def plot_1d_dist(x, draws, injected = None, samples = None, out_folder = '.', name = 'density', label = None, unit = None, show = False, save = True, subfolder = False, true_value = None, true_value_label = '\mathrm{True\ value}', injected_label = '\mathrm{Simulated}', median_label = '\mathrm{Median}', logx = False, logy = False, colors = ['steelblue','darkturquoise','mediumturquoise'], fig = None):
"""
Plot a 1D distribution along with samples from the true distribution (if available).
Differently from plot_median_cr, this method requires the distribution to be already evaluated.
For FIGARO mixture instances, please use plot_median_cr.
Arguments:
iterable x: values at which realisations are evaluated
iterable draws: container for realisations
callable or np.ndarray injected: injected distribution (if available)
np.ndarray samples: samples from the true distribution (if available)
str or Path out_folder: output folder
str name: name to be given to outputs
str label: LaTeX-style quantity label, for plotting purposes
str unit: LaTeX-style quantity unit, for plotting purposes
bool save: whether to save the plots or not
bool show: whether to show the plots during the run or not
bool subfolder: whether to save the plots in different subfolders (for multiple events)
float true_value: true value to infer
str true_value_label: label to assign to the true value marker
str injected_label: label to assign to the injected distribution
str median_label: label to assign to the median distribution
bool logx: x log scale
bool logy: y log scale
list-of-str colors: list of colors for median, 68% and 90% credible regions
matplotlib.figure.Figure fig: figure to use for plotting.
"""
if not np.shape(x)[0] == np.shape(draws)[-1]:
raise ValueError("x and each draw must have the same length")
percentiles = [50, 5, 16, 84, 95]
p = {}
for perc in percentiles:
p[perc] = np.percentile(draws, perc, axis = 0)
color_med, color_68, color_90 = colors
if fig is None:
fig, ax = plt.subplots()
else:
ax = fig.axes[0]
# Samples (if available)
if samples is not None:
ax.hist(samples, bins = int(np.sqrt(len(samples))), histtype = 'step', density = True, label = '$\mathrm{Samples}$')
xlim = ax.get_xlim()
ylim = ax.get_ylim()
# CR
ax.fill_between(x, p[95], p[5], color = color_90, alpha = 0.25)
ax.fill_between(x, p[84], p[16], color = color_68, alpha = 0.25)
# Injection (if available)
if injected is not None:
if callable(injected):
p_x = injected(x)
else:
p_x = injected
if injected_label is not None:
injected_label = '$'+injected_label+'$'
ax.plot(x, p_x, lw = 0.5, color = 'red', label = injected_label)
# Median
if true_value is not None:
if true_value_label is not None:
true_value_label = '$'+true_value_label+'$'
ax.axvline(true_value, ls = '--', color = 'firebrick', lw = 0.5, dashes = (5,5), label = true_value_label)
if median_label is not None:
median_label = '$'+median_label+'$'
ax.plot(x, p[50], lw = 0.7, color = color_med, label = median_label)
if label is None:
label = 'x'
if unit is None or unit == '':
ax.set_xlabel('${0}$'.format(label))
else:
ax.set_xlabel('${0}\ [{1}]$'.format(label, unit))
ax.set_ylabel('$p({0})$'.format(label))
if samples is not None:
ax.set_xlim(xlim)
ax.set_ylim(bottom = 1e-5, top = np.max(p[95])*1.1)
ax.legend(loc = 0)
if logy:
ax.set_yscale('log')
if logx:
ax.set_xscale('log')
fig.align_labels()
if save:
if subfolder:
plot_folder = Path(out_folder, 'density')
if not plot_folder.exists():
try:
plot_folder.mkdir()
except FileExistsError:
# Avoids issue with parallelisation
pass
txt_folder = Path(out_folder, 'txt')
if not txt_folder.exists():
try:
txt_folder.mkdir()
except FileExistsError:
pass
else:
plot_folder = out_folder
txt_folder = out_folder
ax.autoscale(True)
if samples is not None:
ax.set_xlim(xlim)
fig.savefig(Path(plot_folder, '{0}.pdf'.format(name)), bbox_inches = 'tight')
np.savetxt(Path(txt_folder, 'prob_{0}.txt'.format(name)), np.array([x, p[50], p[5], p[16], p[84], p[95]]).T, header = 'x 50 5 16 84 95')
if show:
ax.autoscale(True)
if samples is not None:
ax.set_xlim(xlim)
plt.show()
if (save or show):
plt.close()
return fig
[docs]
def plot_n_clusters_alpha(n_cl, alpha, out_folder = '.', name = 'event', show = False, save = True):
"""
Plot the number of clusters and the concentration parameter as functions of the number of samples.
Arguments:
np.ndarray n_cl: number of active clusters
np.ndarray alpha: concentration parameter
str or Path out_folder: output folder
str name: name to be given to outputs
bool save: whether to save the plot or not
bool show: whether to show the plot during the run or not
"""
fig, ax = plt.subplots()
ax1 = ax.twinx()
ax.plot(np.arange(1, len(n_cl)+1), n_cl, ls = '--', marker = '', lw = 0.7, color = 'k')
ax1.plot(np.arange(1, len(alpha)+1), alpha, ls = '--', marker = '', lw = 0.7, color = 'firebrick')
ax.set_xlabel('$t$')
ax.set_ylabel('$N_{\mathrm{cl}}(t)$', color = 'k')
ax1.set_ylabel('$\alpha(t)$', color = 'firebrick')
if show:
plt.show()
if save:
fig.savefig(Path(out_folder, '{0}_n_cl_alpha.pdf'.format(name)), bbox_inches = 'tight')
plt.close()
[docs]
def pp_plot_cdf(draws, injection, n_points = 1000, out_folder = '.', name = 'event', show = False, save = True):
"""
Make pp-plot comparing draws cdfs and injection cdf
Arguments:
iterable draws: container of mixture instances
callable injection: injected density
int n_points: number of points for linspace
str or Path out_folder: output folder
str name: name to be given to outputs
bool save: whether to save the plot or not
bool show: whether to show the plot during the run or not
"""
all_bounds = np.atleast_2d([d.bounds[0] for d in draws])
x_min = np.max(all_bounds[:,0])
x_max = np.min(all_bounds[:,1])
x = np.linspace(x_min, x_max, n_points)
functions = np.array([mix(x) for mix in draws])
median = np.percentile(functions, 50, axis = 0)
cdf_draws = np.array([fast_cumulative(d) for d in functions])
cdf_median = fast_cumulative(median)
cdf_injection = fast_cumulative(injection(x))
fig = plt.figure()
ax = fig.add_subplot(111, projection = 'pp_plot')
ax.add_confidence_band(len(cdf_median), cl=0.9, color = 'ghostwhite')
ax.add_diagonal()
for cdf in cdf_draws:
ax.plot(cdf_injection, cdf, lw = 0.5, alpha = 0.5, color = 'darkturquoise')
ax.plot(cdf_injection, cdf, color = 'steelblue', lw = 0.7)
ax.set_xlabel('$\mathrm{Injected}$')
ax.set_ylabel('$\mathrm{FIGARO}$')
if show:
plt.show()
if save:
fig.savefig(Path(out_folder, '{0}_ppplot.pdf'.format(name)), bbox_inches = 'tight')
plt.close()
[docs]
def pp_plot_levels(CR_levels, median_CR = None, out_folder = '.', name = 'MDC', show = False, save = True, fig = None, color = 'steelblue'):
"""
Make pp-plot.
Arguments:
iterable CR: 2D array with credible levels for each event
iterable median_CR: credible levels of medians
str or Path out_folder: output folder
str name: name to be given to outputs
bool save: whether to save the plot or not
bool show: whether to show the plot during the run or not
matplotlib.figure.Figure fig: figure to use for plotting.
str color: color for the main line
"""
CR_levels = np.atleast_1d(CR_levels)
if len(CR_levels.shape) > 1:
CR_levels = CR_levels.T
n_evs = CR_levels.shape[-1]
L = np.linspace(0,1,n_evs+2)
if fig is None:
fig = plt.figure()
ax = fig.add_subplot(111, projection = 'pp_plot')
else:
ax = fig.axes[0]
ax.add_confidence_band(n_evs, zorder = n_evs)
ax.add_diagonal(zorder = n_evs+1)
if len(CR_levels.shape) > 1:
sorted = []
for cr in CR_levels:
if median_CR is not None:
lw = 0.3
c = 'lightsteelblue'
else:
lw = 0.6
c = color
x = np.append(0, np.append(cr, 1))
ax.plot(np.sort(x), L, lw = lw, alpha = 0.5, color = c)
if median_CR is not None:
x = np.append(0, np.append(median_CR, 1))
ax.plot(np.sort(x), L, lw = 0.8, color = color, label = '$\mathrm{Median}$', zorder = n_evs+2)
# Add label for draws
handles, labels = ax.get_legend_handles_labels()
line = Line2D([0], [0], label='$\mathrm{Draws}$', lw = lw, color = c)
handles.extend([line])
ax.legend(handles=handles, loc = 0, frameon = False)
else:
x = np.append(0, np.append(CR_levels, 1))
ax.plot(np.sort(x), L, lw = 0.8, color = color, zorder = n_evs+2)
# Maquillage
ax.set_xlabel('$P$')
ax.set_ylabel('$\mathrm{Fraction\ of\ events\ within\ }CR_P$')
if show:
plt.show()
if save:
fig.savefig(Path(out_folder, '{0}_ppplot.pdf'.format(name)), bbox_inches = 'tight')
if (save or show):
plt.close()
return fig
[docs]
def joyplot(draws, x_values, y_values, credible_regions = False, fill = True, solid = False, overlap = 1., xlabel = None, ylabel = None, xunit = None, yunit = None, colormap = 'coolwarm', out_folder = '.', name = 'joyplot', subfolder = False, show = False, save = True, joy = False):
"""
Make a joyplot (also known as ridgeline plot) of a set of distributions.
Heavily inspired by leotac's JoyPy (https://github.com/leotac/joypy).
Arguments:
np.ndarray draws: (n_y, n_x) or (n_y, n_draws, n_x) array containing the distributions to plot
np.ndarray x_values: x values at which the distributions are evaluated
np.ndarray y_values: y values corresponding to the different distributions
bool credible_regions: whether to plot credible regions or not
bool fill: whether to fill the space below the distribution or not (ignored if credible_regions is True)
bool solid: wheter to fill with transparent or solid colour
float overlap: space between different plots
str xlabel: x axis label (LaTeX style)
str ylabel: colorbar label (LaTeX style)
str xunit: x axis units (LaTeX style) - default: no units
str yunit: colorbar units (LaTeX style) - default: no units
str colormap: a valid matplotlib colormap name
str or Path out_folder: output folder
str name: name to be given to outputs
bool subfolder: whether to save the plots in different subfolders (for multiple events)
bool save: whether to save the plots or not
bool show: whether to show the plots during the run or not
bool joy: format the plot to look like the cover of Joy Division's "Unknown Pleasures"
Returns:
matplotlib.figure.Figure: figure with the plot
"""
# Labels
if xlabel is None:
xlabel = '$x$'
else:
xlabel = '${0}$'.format(xlabel)
if xunit is not None:
xlabel = xlabel[:-1]+'\ [{0}]$'.format(xunit)
if ylabel is None:
ylabel = '$y$'
else:
ylabel = '${0}$'.format(ylabel)
if yunit is not None:
ylabel = ylabel[:-1]+'\ [{0}]$'.format(yunit)
# Check that we have as many y values as different draws evaluations
if not len(y_values) == np.shape(draws)[0]:
raise ValueError("y_values and draws must have the same length")
color = iter(colormaps[colormap](np.linspace(1, 0, len(y_values))))
cmappable = plt.cm.ScalarMappable(norm=Normalize(np.min(y_values),np.max(y_values)), cmap=colormap)
num_axes = len(y_values)
# Each plot will have its own axis
fig = plt.figure()
if joy:
fig.patch.set_facecolor('k')
gs = GridSpec(len(y_values), 17, figure = fig)
_axes = []
[_axes.append(fig.add_subplot(gs[i, :-1])) for i in range(len(y_values))]
if not joy:
_axes.append(fig.add_subplot(gs[:, -1]))
# Global colorbar
cbar = fig.colorbar(cmappable, cax = _axes[-1])
cbar.set_label(ylabel)
if joy:
[ax.set_facecolor('k') for ax in _axes]
for i, dd in enumerate(draws[::-1]):
ax = _axes[i]
if not joy:
c = next(color)
else:
c = 'w'
percentiles = [50, 16, 84]
p = {}
if len(np.shape(draws)) == 3:
for perc in percentiles:
p[perc] = np.percentile(dd, perc, axis = 0)
if credible_regions and not joy:
ax.fill_between(x_values, p[84], p[16], color = c, alpha = 0.25)
else:
p[50] = dd
if fill and not credible_regions:
if solid:
alpha = 1
else:
alpha = 0.25
if not joy:
ax.fill_between(x_values, p[50], np.zeros(len(x_values)), color = c, alpha = alpha)
else:
ax.fill_between(x_values, p[50], np.zeros(len(x_values)), color = 'k', alpha = 1)
ax.plot(x_values, p[50], clip_on = True, c = c)
# Setup the current axis: transparency, labels, spines.
ax.yaxis.grid(False)
ax.patch.set_alpha(0)
ax.tick_params(axis='both', which='both', length=0, pad=10)
ax.xaxis.set_visible(False)
ax.yaxis.set_visible(False)
ax.set_frame_on(False)
# Final adjustments
# Compute a final axis, used to apply global settings
last_axis = fig.add_subplot(gs[:, :-1])
if joy:
last_axis.set_facecolor('k')
for side in ['top', 'bottom', 'left', 'right']:
last_axis.spines[side].set_visible(False)
# This looks hacky, but all the axes share the x-axis,
# so they have the same lims and ticks
last_axis.set_xlim(_axes[0].get_xlim())
last_axis.set_xticks(np.array(_axes[0].get_xticks()[1:-1]))
for t in last_axis.get_xticklabels():
if not joy:
t.set_visible(True)
else:
t.set_visible(False)
else:
if not joy:
last_axis.xaxis.set_visible(True)
else:
last_axis.xaxis.set_visible(False)
last_axis.yaxis.set_visible(False)
last_axis.grid(False)
# Last axis on the back
last_axis.zorder = min(a.zorder for a in _axes) - 1
_axes = list(_axes) + [last_axis]
last_axis.set_xlabel(xlabel)
# The magic overlap happens here.
h_pad = 5 + (-5*(1+overlap))
fig.tight_layout(h_pad=h_pad)
if show:
plt.show()
if save:
if not subfolder:
fig.savefig(Path(out_folder, '{0}.pdf'.format(name)), bbox_inches = 'tight')
else:
if not Path(out_folder, 'density').exists():
try:
Path(out_folder, 'density').mkdir()
except FileExistsError:
pass
fig.savefig(Path(out_folder, 'density', '{0}.pdf'.format(name)), bbox_inches = 'tight')
plt.close()
return fig