"""
This module provides some plotting routines for plotting the marginalized posterior
distribution and monitoring performance of the algorithm.
(Other plots are in methods of some classes, e.g. Progress contains the timings plot.)
Under normal circumstances you shouldn't have to use any of the methods in here if you use
the :class:`~gpry.run.Runner` class to run GPry.
"""
import shutil
import warnings
from typing import Sequence, Mapping
from numbers import Number
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
from matplotlib import cm
from tqdm import tqdm # type: ignore
from gpry.surrogate import SurrogateModel
from gpry.mc import process_gdsamples, _name_loglike, _name_logp
from gpry.tools import (
credibility_of_nstd,
nstd_of_1d_nstd,
volume_sphere,
gaussian_distance,
delta_logp_of_1d_nstd,
generic_params_names,
)
# Use LaTeX for labels when a LaTeX toolchain is available; otherwise fall back to
# matplotlib's built-in mathtext (which needs no LaTeX install).
def _latex_available():
# matplotlib's usetex needs a LaTeX engine and dvipng (and, in practice, ghostscript)
return all(shutil.which(exe) for exe in ("latex", "dvipng"))
plt.rcParams["text.usetex"] = _latex_available()
_plot_dist_fontsize = 7
[docs]
def simple_latex_sci_notation(string):
"""
If ``string`` contains a ``%g`` or ``%e`` number representation, substitutes the ``e``
for a power of 10.
It does *not* add dollars around the string.
NB: it assumes that the string passed contains a single number, and nothing else.
"""
if "e" not in string:
return string
sigfigs, exp = string.split("e")
sign = "" if exp.startswith("+") else "-"
return f"{sigfigs}\\cdot 10^{{{sign}{exp[1:].lstrip('0')}}}"
[docs]
def param_samples_for_slices(X, i, bounds, n=200):
"""
From an array of points `X = [X^i] = [[X^1_1, X^1_2,...], [X^2_1,...], ...]`, it
generates a list of points per sample, where the `i` coordinate is sliced within the
region defined by `bounds`, and the rest of them are kept fixed.
"""
X = np.atleast_2d(X)
Xs_i = np.linspace(bounds[0], bounds[1], n)
X_slices = np.empty(shape=(X.shape[0], n, X.shape[1]), dtype=float)
for j, X_j in enumerate(X):
X_slices[j, :, :] = n * [X_j]
X_slices[:, :, i] = Xs_i
return X_slices
[docs]
def prepare_slices_func(func, X_fiducial, bounds, indices=None, n=50):
"""
Prepare slices of the given function,
Parameters
----------
func : callable
Function for which to prepare slices. It needs to take arguments in the way
``X_fiducial`` is passed: ``func(*X_fiducial)`` if ``X_fiducual`` is a list, or
``func(**X_fiducial)`` if X_fiducial is a dictionary.
X_fiducial : array-like, shape = (n_dimensions), or dict
Fiducial point for the slices: slice ``i`` corresponds to fixing all parameters
but that with index ``i``, which is evaluaded on a grid within its bounds. It can
be a dictionary with arguments of ``func`` as keys.
bounds : array-like, shape = (n_dimensions, 2), or dict
Bounds for the slices per parameter.
indices : list, optional
A list of integers (if ``X_fiducial`` is a list) or parameter names (if
``X_fiducial`` is a dict), denoting the parameters for which the slices will be
prepared (all of them, if left unspecified).
n : int
Number of samples per slice (default: 50). Careful if the posterior is slow!
Returns
-------
indices, params, Xs, ys : list(int), list(str) len=len(indices), array-like
shape=(len(indices), n, dim)), array-like shape=(len(indices, n))
"""
# Parse and check input
if isinstance(X_fiducial, Mapping):
is_kwarg = True
dim = len(X_fiducial)
X_fiducial = np.array(list(X_fiducial.values()))
params = list(X_fiducial)
if indices is None:
indices = params
try: # Assumes indices is a list of str
indices = [params.index(p) for p in indices]
except ValueError as excpt:
raise ValueError(
"`indices` is not a list of parameter names, or contains names not in "
"`X_fiducial`."
) from excpt
try: # Assumes bonds is a Mapping
bounds = [bounds[p] for p in params]
except (TypeError, IndexError) as excpt:
raise ValueError(
"`bounds` is not a dict, or bounds could not be founds for all "
"parameters."
) from excpt
else:
is_kwarg = False
X_fiducial = np.atleast_1d(X_fiducial)
dim = len(X_fiducial)
params = generic_params_names(dim)
if indices is None:
indices = list(range(dim))
# Assumes indices is a list of int
if not all((isinstance(i, int) and i < dim) for i in indices):
raise ValueError(
"`indices` is not a list of integer indices, or contains indices larger "
"than the length of `X_fiducial`."
)
# Assumes bonds is a (2d) list
if len(bounds) < len(params):
raise ValueError("`bounds` is not a list of bounds of the right lenght.")
if not isinstance(n, int) or n < 2:
raise ValueError("`n` must be a positive integer > 2.")
# Prepare and evaluate slices
Xs, ys = np.empty(shape=(len(indices), n, dim)), np.empty(shape=(len(indices), n))
for j, index in enumerate(indices):
Xs[j] = param_samples_for_slices([X_fiducial], index, bounds[index], n=n)[0]
progress_bar_desc = f"Slicing param {j + 1} of {len(indices)}"
for k, x in tqdm(enumerate(Xs[j]), total=n, desc=progress_bar_desc):
try:
if is_kwarg:
x_arg = dict(zip(params, x))
ys[j][k] = func(**x_arg)
else:
x_arg = x
ys[j][k] = func(*x_arg)
except TypeError as excpt:
raise TypeError(
f"Could not call the target function with arguments {x_arg}. Maybe "
"`X_fiducial` contained keys that are not argument names of the "
f"function? Err msg: {excpt}"
) from excpt
except Exception as excpt:
warnings.warn(
f"The function failed when called with arguments {x_arg}. Using NaN."
f"Err masg: {excpt}"
)
ys[j][k] = np.nan
return indices, [params[i] for i in indices], Xs, ys
# NB: use this one in the future to reformulate the other ones: more generic
[docs]
def plot_slices_func(
func,
X_fiducial,
bounds,
indices=None,
n=50,
fig_kwargs=None,
labels=None,
):
"""
Plot slices of the given function,
Parameters
----------
func : callable
Function for which to prepare slices. It needs to take arguments in the way
``X_fiducial`` is passed: ``func(*X_fiducial)`` if ``X_fiducual`` is a list, or
``func(**X_fiducial)`` if X_fiducial is a dictionary.
X_fiducial : array-like, shape = (n_dimensions), or dict
Fiducial point for the slices: slice ``i`` corresponds to fixing all parameters
but that with index ``i``, which is evaluaded on a grid within its bounds. It can
be a dictionary with arguments of ``func`` as keys.
bounds : array-like, shape = (n_dimensions, 2), or dict
Bounds for the slices per parameter.
indices : list, optional
A list of integers (if ``X_fiducial`` is a list) or parameter names (if
``X_fiducial`` is a dict), denoting the parameters for which the slices will be
prepared (all of them, if left unspecified).
n : int
Number of samples per slice (default: 50). Careful if the posterior is slow!
fig_kwargs : dict, optional
Dict of kw arguments to pass to the `subplots` constructor. Only ``layout``,
``dpi`` considered safe.
labels : lst(str), optional
Strings (possibly Latex) to use for axes labels. Length cases: None or len=0:
plain parameter names for x labels and no y label; len=1: used as y label, plain
names for x labels; len=len(indices): used as x labels; len=len(indices)+1: used
as x labels and y label, in that order.
Returns
-------
fig, axarr: figure and array of axes used for the plot.
"""
indices, params, Xs, ys = prepare_slices_func(
func, X_fiducial, bounds, indices=indices, n=n
)
if not isinstance(labels, Sequence) or len(labels) == 0:
x_labels = params
y_label = None
elif len(labels) == 1:
x_labels = params
y_label = labels[0]
elif len(labels) == len(params):
x_labels = labels
y_label = None
elif len(labels) == len(params) + 1:
x_labels = labels[:-1]
y_label = labels[-1]
else:
raise ValueError("Value for `labels` not recognised, or length not valid.")
fig_kwargs_defaults = {
"nrows": 1,
"ncols": len(indices),
"layout": "constrained",
"figsize": (4 * len(indices), 2),
"dpi": 200,
}
fig_kwargs_defaults.update(fig_kwargs or {})
fig, axes = plt.subplots(**fig_kwargs_defaults)
if not isinstance(axes, Sequence):
axes = [axes]
color = "tab:blue"
for j, (i, p) in enumerate(zip(indices, params)):
axes[j].axvline(X_fiducial[i], c="0.75", ls="--")
axes[j].plot(Xs[j, :, i], ys[j], c=color)
axes[j].scatter(Xs[j, :, i], ys[j], marker=".", s=10, c=color)
axes[j].set_xlabel(x_labels[j])
axes[j].set_ylabel(y_label)
return fig, axes
[docs]
def plot_slices(truth, surrogate, acquisition, X=None, reference=None):
"""
Plots slices along parameter coordinates for a series `X` of given points (the
surrogate's training set if not specified). For each coordinate, there is a slice per
point, leaving all coordinates of that point fixed except for the one being sliced.
Lines are coloured according to the value of the mean GP at points X.
# TODO: make acq func optional
"""
params = truth.params
fig, axes = plt.subplots(
nrows=2,
ncols=len(params),
sharex="col",
layout="constrained",
figsize=(4 * len(params), 4),
dpi=200,
)
# Define X to plot
if X is None:
X = surrogate.X_regress
y = surrogate.y_regress
else:
y = surrogate.predict(X)
min_y, max_y = min(y), max(y)
norm_y = lambda y: (y - min_y) / (max_y - min_y)
prior_bounds = truth.prior_bounds
Xs_for_plots = dict(
(p, param_samples_for_slices(X, i, prior_bounds[i], n=200))
for i, p in enumerate(params)
)
if reference is not None:
reference = _prepare_reference(reference, truth)
cmap = matplotlib.colormaps["viridis"]
for i, p in enumerate(params):
for j, Xs_j in enumerate(Xs_for_plots[p]):
cmap_norm = cmap(norm_y(y[j]))
alpha = 1
# TODO: could cut by half # of GP evals by reusing for acq func
axes[0, i].plot(
Xs_j[:, i], surrogate.predict(Xs_j), c=cmap_norm, alpha=alpha
)
axes[0, i].scatter(X[j][i], y[j], color=cmap_norm, alpha=alpha)
axes[0, i].set_ylabel(r"$\log(p)$")
acq_values = acquisition(Xs_j, surrogate)
axes[1, i].plot(Xs_j[:, i], acq_values, c=cmap_norm, alpha=alpha)
axes[1, i].set_ylabel(r"$\alpha(\mu,\sigma)$")
label = truth.labels[i] if truth.labels is not None else p
if label != p:
label = "$" + label + "$"
axes[1, i].set_xlabel(label)
bounds = (reference or {}).get(p)
if bounds is not None:
for ax in axes[:, i]:
if len(bounds) == 5:
ax.axvspan(
bounds[0],
bounds[4],
facecolor="tab:blue",
alpha=0.2,
zorder=-99,
)
ax.axvspan(
bounds[1],
bounds[3],
facecolor="tab:blue",
alpha=0.2,
zorder=-99,
)
ax.axvline(bounds[2], c="tab:blue", alpha=0.3, ls="--")
[docs]
def plot_slices_reference(truth, surrogate, X, plot_truth=True, reference=None):
"""
Plots slices of the surrogate model and true log-posterior (if ``plot_truth=True``)
along parameter coordinates for a given point ``X``, leaving all coordinates of that
point fixed except for the one being sliced.
"""
params = truth.params
fig, axes = plt.subplots(
nrows=1,
ncols=len(params),
sharex="col",
layout="constrained",
figsize=(4 * len(params), 2),
dpi=200,
)
prior_bounds = truth.prior_bounds
if X is None:
if reference is None:
raise ValueError("Needs at least a reference point or a reference sample.")
# TODO: if reference given as a sample, take best point from it.
X_array = np.array([X[p] for p in params])
# y_surrogate_centre = surrogate.predict(np.atleast_2d(X_array))[0]
if plot_truth:
y_truth_centre = truth.logp(X_array)
(
Xs_for_plots,
ys_surrogate_for_plot,
sigmas_surrogate_for_plot,
ys_truth_for_plot,
) = ({}, {}, {}, {})
for i, p in enumerate(params):
Xs_for_plots[p] = param_samples_for_slices(
[X_array], i, prior_bounds[i], n=200
)[0]
ys_surrogate_for_plot[p], sigmas_surrogate_for_plot[p] = surrogate.predict(
Xs_for_plots[p], return_std=True
)
if plot_truth:
ys_truth_for_plot[p] = np.array([truth.logp(x) for x in Xs_for_plots[p]])
# Training set referenced to X and div by X, for distance-based transparency
X_train_diff = (surrogate.X_regress - X_array) / X_array
if reference is not None:
reference = _prepare_reference(reference, truth)
for i, p in enumerate(params):
axes[i].fill_between(
Xs_for_plots[p][:, i],
ys_surrogate_for_plot[p] - 2 * sigmas_surrogate_for_plot[p],
ys_surrogate_for_plot[p] + 2 * sigmas_surrogate_for_plot[p],
color="0.5",
alpha=0.25,
edgecolor="none",
)
axes[i].fill_between(
Xs_for_plots[p][:, i],
ys_surrogate_for_plot[p] - 1 * sigmas_surrogate_for_plot[p],
ys_surrogate_for_plot[p] + 1 * sigmas_surrogate_for_plot[p],
color="0.5",
alpha=0.25,
edgecolor="none",
)
if plot_truth:
axes[i].plot(Xs_for_plots[p][:, i], ys_truth_for_plot[p], ls="--")
axes[i].scatter(X[p], y_truth_centre, marker="*")
axes[i].set_ylabel(r"$\log(p)$")
label = truth.labels[i] if truth.labels is not None else p
if label != p:
label = "$" + label + "$"
axes[i].set_xlabel(label)
# If there is an infinities classifier, use if for the lower bound on y:
diff_min_logp = getattr(surrogate, "diff_threshold", None)
if diff_min_logp is not None:
try:
max_y = max(ys_surrogate_for_plot[p])
upper_y = max_y
if plot_truth:
upper_y = max(max_y, *ys_truth_for_plot[p])
axes[i].set_ylim(
max_y - 1.05 * diff_min_logp, upper_y + 0.05 * diff_min_logp
)
except ValueError as e:
print(
f"ERROR when setting y-lims for '{p}': max(y) was "
f"{max(ys_surrogate_for_plot[p])}, "
f"diff_threshold was {diff_min_logp}, lower_bound was "
f"{max(ys_surrogate_for_plot[p]) - 1.05 * diff_min_logp}, "
f"upper bound was {upper_y}, MSG was {e}"
)
# Add training set
dists = np.sqrt(
np.sum(np.power(np.delete(X_train_diff, i, axis=-1), 2), axis=-1)
)
dists_relative = dists / max(dists)
axes[i].scatter(
surrogate.X_regress[:, i],
surrogate.y_regress,
marker=".",
alpha=1 - dists_relative,
zorder=-9,
)
bounds = (reference or {}).get(p)
if bounds is not None:
if len(bounds) == 5:
axes[i].axvspan(
bounds[0], bounds[4], facecolor="tab:blue", alpha=0.2, zorder=-99
)
axes[i].axvspan(
bounds[1], bounds[3], facecolor="tab:blue", alpha=0.2, zorder=-99
)
axes[i].axvline(bounds[2], c="tab:blue", alpha=0.3, ls="--")
[docs]
def force_agg_backend(func):
"""
Temporarily switch to Agg backend -- solves getdist/matplotlib mem leak.
"""
def wrapper(*args, **kwargs):
prev_backend = matplotlib.get_backend()
matplotlib.use("Agg")
try:
return func(*args, **kwargs)
except Exception as excpt:
raise excpt
finally:
# Switch back to prev backend
matplotlib.use(prev_backend)
return wrapper
[docs]
@force_agg_backend
def plot_corner_getdist(
mc_samples,
params=None,
bounds=None,
filled=None,
training=None,
training_highlight_last=False,
add_logp=False,
add_loglike=False,
markers=None,
output=None,
output_dpi=200,
subplot_size=2,
close=False,
):
"""
Creates a corner plot the given MC samples, and optionally shows evaluation locations.
If called repeatedly, it may leak memory, unfortunately. To avoid this, execute it
like this:
.. code:: python
# Temporarily switch to Agg backend
prev_backend = matplotlib.get_backend()
matplotlib.use("Agg")
try:
plot_corner_getdist(...)
except:
...
finally:
# Switch back to prev backend
mpl.use(prev_backend)
Parameters
----------
mc_samples: dict(str, (cobaya.SampleCollection, getdist.MCSamples, str))
Dict of MC samples, with their plot label as key, and the sample as value, either
as GetDist or Cobaya samples, or as a path where there are samples saved.
params : list(str), optional
List of parameter names to be plotted, by default all of the ones in the first
MC sample, included derived ones like probability densities.
bounds : array-like, shape = (len(params), 2), list(array-like, shape = (2))
Dict or list (sorted as ``params``) of parameter bounds.
filled : dict(str, bool), list
Dictionary with labels as keys specifying the `filled` property of the contours.
Contours are filled by default when unspecified (including key missing for a
passed sample). If it is a list, the same order as in ``mc_samples`` is assumed.
training : SurrogateModel, dict(str or tuple, SurrogateModel), optional
If a surrogate model is passed, it plots the training samples (including the
discarded ones) on top of the contours. Samples outside the axes ranges are not
plotted.
The parameters of the surrogate model need to be assumed, since the surrogate
model does not save names: if a surrogate is passed, the sampled parameters of the
first MC sample will be used; if a single-key dict is passed with a str as a key,
it will used the parameter names from the MC sample with that label; if the key is
a tuple of strings, they will be used as parameters
training_highlight_last : bool (default: False)
If ``training`` is True, whether the last-added training points are highlighted.
add_logp : bool (default: False)
Whether to add a row with the log-posterior as a parameter (if available as a
derived parameter of the last input samples. Incompatible with ``add_loglike``.
add_loglike : bool (default: False)
Whether to add a row with the log-likelihood as a parameter (if available as a
derived parameter of the last input samples. Incompatible with ``add_logp``.
subplot_size : float, default = 2
Size of each subplot in the corner plot.
output : str, optional (default=None)
Path, including name and extension, of the saved figure.
Not saved if left unspecified.
output_dpi : int (default: 200)
The resolution of the generated plot in DPI.
close : bool (default: False)
Whether to close the figure (it cannot be shown with ``plt.show()`` after calling
this function). It should be True if called multiple times during a run, to avoid
wasting memory.
Returns
-------
getdist.plots.GetDistPlotter object containing the figure.
"""
if close and output is None:
raise ValueError("Calling with 'close=True' and no output has no effect.")
if add_logp and add_loglike:
raise ValueError("Only one of add_logp/add_loglike can be True.")
if _name_logp in params or _name_loglike in params:
raise ValueError(
"To plot log-posterior or log-likelihood, do not add them as parameters; "
"use add_logp|add_loglike instead."
)
if not isinstance(mc_samples, Mapping):
raise TypeError(
"The first argument must be a list of MC samples with the sample legend "
"labels as keys."
)
gdsamples_dict = process_gdsamples(mc_samples)
# Prepare training samples early -- fail asap.
training_params = None
if training is not None:
if isinstance(training, Mapping):
training_key = list(training)[0]
training = list(training.values())[0]
if isinstance(training_key, tuple):
training_params = training_key
else: # assumed key corresponding to passed mcsamples
if training_key not in mc_samples:
raise ValueError(
"`training` passed as dict, but key not found in mc_samples."
)
training_params = (
gdsamples_dict[training_key].getParamNames().getRunningNames()
)
elif isinstance(training, SurrogateModel):
# Use first MC passed
training_params = (
list(gdsamples_dict.values())[0].getParamNames().getRunningNames()
)
else:
raise TypeError("'training' is not a SurrogateModel instance.")
import getdist.plots as gdplt # type: ignore
gdplot = gdplt.get_subplot_plotter(subplot_size=subplot_size, auto_close=close)
gdplot.settings.line_styles = "tab10"
gdplot.settings.solid_colors = "tab10"
triang_args = [list(gdsamples_dict.values())]
if params is not None:
# GetDist failsafe: can only plot the params of the last sample in the input
last_samples = list(gdsamples_dict.values())[-1]
at_most_params = set(last_samples.getParamNames().getRunningNames()).union(
set(last_samples.getParamNames().getDerivedNames())
)
if add_logp:
params = list(params) + [_name_logp]
elif add_loglike:
params = list(params) + [_name_loglike]
params = [p for p in params if p in at_most_params]
triang_args.append(params)
if isinstance(bounds, Mapping):
bounds = {p: bounds.get(p) for p in params}
elif isinstance(bounds, (Sequence, np.ndarray)):
if len(bounds) != len(params):
raise ValueError("`bounds` and `params` have different number of elements.")
bounds = {p: bounds[i] for i, p in enumerate(params)}
if isinstance(filled, Sequence):
# Assume it refers to the first samples if not complete; the rest are filled.
filled = {
name: (filled[i] if i < len(filled) else True)
for i, name in enumerate(gdsamples_dict)
}
triang_kwargs = {
"legend_labels": list(gdsamples_dict),
"filled": [(filled or {}).get(k, True) for k in gdsamples_dict],
"param_limits": bounds or {},
"markers": markers,
}
try:
gdplot.triangle_plot(*triang_args, **triang_kwargs)
except Exception as excpt:
raise ValueError(
f"Could not do corner plot. GetDist err. msg.: {excpt}"
) from excpt
if training is not None and training.d > 1:
getdist_add_training(
gdplot, training_params, training, highlight_last=training_highlight_last
)
if output is not None:
plt.savefig(output, dpi=output_dpi)
return gdplot
[docs]
def getdist_add_training(
getdist_plot,
params,
surrogate,
colormap="viridis",
marker=".",
marker_inf="x",
highlight_last=False,
):
"""
Adds the training points to a GetDist triangle plot, coloured according to
their log-posterior value.
Parameters
----------
getdist_plot : `GetDist triangle plot <https://getdist.readthedocs.io/en/latest/plots.html?highlight=triangle_plot#getdist.plots.GetDistPlotter.triangle_plot>`_
Contains the marginalized contours and potentially other things.
params : list(str)
The assumed parameter names for the surrogate samples. Need to be a subset of the
ones plotted by the GetDistPlotter.
surrogate : SurrogateModel
The trained GP Regressor containing the samples.
colormap : matplotlib colormap, optional (default="viridis")
Color map from which to get the color scale to represent the GP model value for
the training points.
marker : matplotlib marker, optional (default=".")
Marker to be used for the training points.
marker_inf : matplotlib marker, optional (default=".")
Marker to be used for the non-finite training points.
highlight_last: bool (default=False)
Draw a red circle around the points added in the last iteration
Returns
-------
The GetDist triangle plot with the added training points.
"""
# Gather axes and bounds
d = len(params)
ax_dict = {}
bounds = [None] * len(params)
for i, pi in enumerate(params):
for j, pj in enumerate(params):
ax = getdist_plot.get_axes_for_params(pi, pj, ordered=True)
if not ax:
continue
ax_dict[(i, j)] = ax
bounds[i] = ax.get_xlim()
bounds[j] = ax.get_ylim()
# Now reduce the set of points to the ones within ranges
# (needed to get good limits for the colorbar of the log-posterior)
Xs_finite = surrogate.X_regress
ys_finite = surrogate.y_regress
Xs_infinite = surrogate.X_infinite
for i, (mini, maxi) in enumerate(bounds):
i_within_finite = np.argwhere(
np.logical_or(mini < Xs_finite[:, i], Xs_finite[:, i] < maxi)
)
Xs_finite = np.atleast_2d(np.squeeze(Xs_finite[i_within_finite]))
ys_finite = np.atleast_1d(np.squeeze(ys_finite[i_within_finite]))
i_within_infinite = np.argwhere(
np.logical_or(mini < Xs_infinite[:, i], Xs_infinite[:, i] < maxi)
)
Xs_infinite = np.atleast_2d(np.squeeze(Xs_infinite[i_within_infinite]))
if highlight_last:
Xs_last = surrogate.X_last_appended
i_within_last = np.argwhere(
np.logical_or(mini < Xs_last[:, i], Xs_last[:, i] < maxi)
)
Xs_last = np.atleast_2d(np.squeeze(Xs_last[i_within_last]))
if (
len(Xs_finite) == 0 and len(Xs_infinite) == 0
): # no points within plotting ranges
return getdist_plot
# Create colormap with appropriate limits
cmap = matplotlib.colormaps[colormap]
if len(Xs_finite):
Ncolors = 256
color_bounds = np.linspace(min(ys_finite), max(ys_finite), Ncolors)
norm = matplotlib.colors.BoundaryNorm(color_bounds, Ncolors)
# Add points
for (i, j), ax in ax_dict.items():
# 1st -inf points, so the are displayed in the background of the finite ones.
# and we give them low zorder anyway, so that they lay behind the contours
if len(Xs_infinite) > 0:
points_infinite = Xs_infinite[:, [i, j]]
ax.scatter(
*points_infinite.T,
marker=marker_inf,
s=20,
c="k",
alpha=0.3,
zorder=-99,
)
if len(Xs_finite) > 0:
points_finite = Xs_finite[:, [i, j]]
ax.scatter(
*points_finite.T, marker=marker, c=norm(ys_finite), alpha=0.3, cmap=cmap
)
if highlight_last and len(Xs_last) > 0:
points_last = Xs_last[:, [i, j]]
ax.scatter(
*points_last.T,
marker="o",
c=len(points_last) * [[0, 0, 0, 0]],
edgecolor="r",
lw=0.5,
)
# Colorbar
if len(Xs_finite) > 0 and not np.isclose(min(ys_finite), max(ys_finite)):
getdist_plot.fig.colorbar(
cm.ScalarMappable(norm=norm, cmap=cmap),
label=r"$\log(p)$",
ax=getdist_plot.fig.add_axes(
[1 - 0.2 / d, 1 - 0.85 / d, 0.5 / d, 0.5 / d],
frame_on=False,
xticks=[],
yticks=[],
),
ticks=np.linspace(min(ys_finite), max(ys_finite), 5),
location="left",
)
return getdist_plot
[docs]
def plot_convergence(
convergence_criterion,
evaluations="total",
marker="",
axes=None,
ax_labels=True,
legend_loc="upper right",
):
"""
Plots the value of the convergence criterion as function of the number of
(accepted) training points.
Parameters
----------
convergence_criterion : The instance of the convergence criterion which has
been called in the BO loop
evaluations : "total" or "accepted"
Whether to plot the total number of posterior evaluations or only the
accepted steps.
marker : matplotlib marker, optional (default="")
Marker used for the plot. Will be passed to ``matplotlib.pyplot.plot``.
axes : matplotlib axes, optional
Axes to be used, if passed.
ax_labels : bool, optional (default: True)
Add axes labels.
legend_loc : str (default: "upper right")
Location of the legend.
Returns
-------
The plot convergence criterion vs. number of training points
"""
if not isinstance(convergence_criterion, Sequence):
convergence_criterion = [convergence_criterion]
if axes is None:
fig, axes = plt.subplots()
else:
fig = axes.get_figure()
for i, cc in enumerate(convergence_criterion):
color = plt.rcParams["axes.prop_cycle"].by_key()["color"][i]
values, n_posterior_evals, n_accepted_evals = cc.get_history()
name = cc.__class__.__name__
n_evals = np.array(
{"total": n_posterior_evals, "accepted": n_accepted_evals}[evaluations],
dtype=int,
)
try:
axes.plot(n_evals, values, marker=marker, color=color, label=name)
except KeyError as excpt:
raise ValueError(
"'evaluations' must be either 'total' or 'accepted'."
) from excpt
if hasattr(cc, "limit"):
axes.axhline(cc.limit, ls="--", lw=1, c=color)
if ax_labels:
axes.set_xlabel(f"{evaluations} number of posterior evaluations")
axes.set_ylabel("Value of convergence criterion")
axes.set_yscale("log")
axes.grid(axis="y")
axes.legend(loc=legend_loc, prop={"size": _plot_dist_fontsize})
return fig, axes
def _prepare_reference(
reference,
truth,
):
"""
Turns `reference` into a dict with parameters as keys and a list of 5 numbers as
values: two lower bounds, a central value, and two upper bounds, e.g. percentiles
5, 25, 50, 75, 95.
If getdist.MCSamples passed, bounds are by default 68% and 95%, and the central value
is the mean.
"""
# Ensure it is a dict
try:
from getdist import MCSamples
if isinstance(reference, MCSamples):
means = reference.getMeans()
margstats = reference.getMargeStats()
bounds = {}
for p in truth.params:
# NB: numerOfName doest not use renames; needs to find "original" name
p_in_ref = reference.paramNames.parWithName(p).name
i_p = reference.paramNames.numberOfName(p_in_ref)
# by default lims/contours are [68, 95, 99]
try:
lims = margstats.parWithName(p).limits
except AttributeError as excpt:
raise ValueError(
f"Could not find parameter {p} in reference sample, which "
f"includes {reference.getParamNames().list()})"
) from excpt
bounds[p] = [
lims[1].lower,
lims[0].lower,
means[i_p],
lims[0].upper,
lims[1].upper,
]
reference = bounds
except ModuleNotFoundError: # getdist not installed
return None
if not isinstance(reference, Mapping):
# Assume parameters in order; check right number of them
if len(reference) != truth.d:
raise ValueError(
"reference must be a list containing bounds per parameter for all of them"
", or a dict with parameters as keys and these same values."
)
reference = dict(zip(truth.params, reference))
# Ensure it contains all parameters and 5 numbers (or None's) per parameter
for p in truth.params:
if p not in reference:
reference[p] = [None] * 5
values = reference[p]
if isinstance(values, Number):
values = [values]
if len(values) == 1:
reference[p] = [None, None] + list(values) + [None, None]
elif len(values) != 5:
raise ValueError(
"the elements of reference must be either a single central value, or a "
"list of 5 elements: [lower_bound_2, lower_bound_1, central_value, "
"upper_bound_2, upper_bound_1]."
)
return reference
[docs]
def plot_trace(
truth,
surrogate,
convergence_criterion,
progress,
colormap="viridis",
reference=None,
):
"""
Plots the evolution of the run along true model evaluations, showing evolution of the
convergence criterion and the values of the log-posterior and the individual
parameters.
Can take a reference sample or reference bounds (dict with parameters as keys and 5
sorted bounds as values, or alternatively just a central value).
"""
X = surrogate.X
y = surrogate.y
if surrogate.infinities_classifier is not None:
y_finite = surrogate.is_finite_y(y)
else:
y_finite = np.full(shape=len(y), fill_value=True)
if reference is not None:
reference = _prepare_reference(reference, truth)
fig, axes = plt.subplots(
nrows=2 + truth.d,
ncols=1,
sharex=True,
layout="constrained",
figsize=(min(4, 0.3 * len(X)), 1.5 * (2 + X.shape[1])),
dpi=400,
)
i_eval = list(range(1, 1 + len(X)))
# TOP: convergence plot
try:
plot_convergence(
convergence_criterion,
evaluations="total",
marker="",
axes=axes[0],
ax_labels=False,
legend_loc="lower left",
)
except ValueError: # no criterion computed yet
pass
axes[0].set_ylabel("Conv. crit.")
# 2nd: posterior plot
kwargs_accepted = {
"marker": ".",
"linewidths": 0.1,
"edgecolor": "0.1",
"cmap": colormap,
}
axes[1].scatter(i_eval, y, c=np.where(y_finite, y, np.inf), **kwargs_accepted)
# Gaussian contours
dashdotdotted = (0, (3, 5, 1, 5, 1, 5))
nsigmas_styles = {1: "-", 2: "--", 5: "-.", 10: ":", 20: dashdotdotted}
y_min_plot, y_max_plot = axes[1].get_ylim()
y_max = np.max(y)
for ns, nsls in nsigmas_styles.items():
y_ns = y_max - delta_logp_of_1d_nstd(ns, truth.d)
if y_ns > y_min_plot:
axes[1].axhline(
y_ns,
ls=nsls,
c="0.3",
lw=0.75,
zorder=-1,
label=f"{ns}-$\\sigma$ (Gauss. approx.)",
)
axes[1].set_ylabel(r"$\log(p)$")
axes[1].grid(axis="y")
axes[1].legend(loc="lower left", prop={"size": _plot_dist_fontsize})
# Kernel scales
output_scale, length_scales = surrogate.scales
scales_kwargs = {
"verticalalignment": "center",
"horizontalalignment": "right",
"fontsize": _plot_dist_fontsize,
"bbox": {
"facecolor": "white",
"alpha": 0.5,
},
}
axes[1].text(
0.965,
0.12,
f"Output scale: ${simple_latex_sci_notation(f'{output_scale:.2g}')}$",
transform=axes[1].transAxes,
**scales_kwargs,
)
# Single parameters plots
for i, p in enumerate(truth.params):
label = truth.labels[i] if truth.labels else p
ax = axes[i + 2]
if surrogate.infinities_classifier is not None and sum(y_finite) < len(X):
ax.scatter(
i_eval,
X[:, i],
marker="x",
c=np.where(y_finite, None, 0.5),
cmap="gray",
vmin=0,
vmax=1,
s=20,
)
ax.scatter(
i_eval,
X[:, i],
c=np.where(y_finite, y, np.inf),
**kwargs_accepted,
)
bounds = (reference or {}).get(p)
if bounds is not None:
if len(bounds) == 5:
ax.axhspan(
bounds[0], bounds[4], facecolor="tab:blue", alpha=0.2, zorder=-99
)
ax.axhspan(
bounds[1], bounds[3], facecolor="tab:blue", alpha=0.2, zorder=-99
)
ax.axhline(bounds[2], c="tab:blue", alpha=0.3, ls="--")
ax.set_ylabel("$" + label + "$" if label != p else p)
ax.grid(axis="y")
# Length scales
ax.text(
0.965,
0.12,
f"Length scale: ${simple_latex_sci_notation(f'{length_scales[i]:.2g}')}$",
transform=ax.transAxes,
**scales_kwargs,
)
# Format common x-axis
axes[0].set_xlim(0, len(X) + 0.5)
axes[-1].set_xlabel("Number of posterior evaluations")
n_train = progress.data["n_total"][1]
for ax in axes:
ax.axvspan(0, n_train + 0.5, facecolor="0.85", zorder=-999)
for n_iteration in progress.data["n_total"][1:]:
ax.axvline(n_iteration + 0.5, ls="--", c="0.75", lw=0.75, zorder=-9)
# TODO: make sure the x ticks are int
[docs]
def plot_distance_distribution(
points, mean, covmat, density=False, show_added=True, ax=None
):
"""
Plots a histogram of the distribution of points with respect to the number of standard
deviations. Confidence level boundaries (Gaussian approximantion, dimension-dependent)
are shown too.
Parameters
----------
points: array-like, with shape ``(N_points, N_dimensions)``, or surrogate instance
Points to be used for the histogram.
mean: array-like, ``(N_dimensions)``.
Mean of the distribution.
covmat: array-like, ``(N_dimensions, N_dimensions)``.
Covariance matrix of the distribution.
density: bool (default: False)
If ``True``, bin height is normalised to the (hyper)volume of the (hyper)spherical
shell corresponding to each standard deviation.
show_added: bool (default True)
Colours the stacks depending on how early or late the corresponding points were
added (bluer stacks represent newer points).
ax: matplotlib axes
If provided, they will be used for the plot.
Returns
-------
Tuple of current figure and axes ``(fig, ax)``.
"""
if isinstance(points, SurrogateModel):
points = points.X_regress
dim = np.atleast_2d(points).shape[1]
radial_distances = gaussian_distance(points, mean, covmat)
bins = list(range(0, int(np.ceil(np.max(radial_distances))) + 1))
num_or_dens = "Density" if density else "Number"
if density:
volumes = [
volume_sphere(bins[i], dim) - volume_sphere(bins[i - 1], dim)
for i in range(1, len(bins))
]
weights = [1 / volumes[int(np.floor(r))] for r in radial_distances]
else:
weights = np.ones(len(radial_distances))
if ax is None:
fig, ax = plt.subplots()
else:
fig = ax.get_figure()
title_str = f"{num_or_dens} of points per standard deviation"
if show_added:
title_str += " (bluer=newer)"
cmap = plt.get_cmap("Spectral")
colors = [cmap(i / len(points)) for i in range(len(points))]
ax.hist(
np.atleast_2d(radial_distances),
bins=bins,
weights=np.atleast_2d(weights),
color=colors,
stacked=True,
)
else:
ax.hist(radial_distances, bins=bins, weights=weights)
ax.set_title(title_str)
# cls = [credibility_of_nstd(s, 1) for s in [1, 2, 3, 4]] # use 1d cl's as reference
nstds = [1, 2, 3, 4]
linestyles = ["-", "--", "-.", ":"]
for nstd, ls in zip(nstds, linestyles):
std_of_cl = nstd_of_1d_nstd(nstd, dim)
if std_of_cl < max(radial_distances):
ax.axvline(
std_of_cl,
c="0.75",
ls=ls,
zorder=-99,
label=f"${100 * credibility_of_nstd(std_of_cl, dim):.2f}\\%$ prob mass",
)
ax.set_ylabel(f"{num_or_dens} of points")
ax.set_xlabel("Number of standard deviations")
ax.legend(loc="upper right")
return (fig, ax)
def _plot_2d_model_acquisition(surrogate, acquisition, last_points=None, res=200):
"""
Contour plots for model prediction and acquisition function value of a 2d model.
If ``last_points`` passed, they are highlighted.
"""
if surrogate.d != 2:
warnings.warn("This plots are only possible in 2d.")
return
# TODO: option to restrict bounds to the min square containing traning samples,
# with some padding
bounds = surrogate.bounds
x = np.linspace(bounds[0][0], bounds[0][1], res)
y = np.linspace(bounds[1][0], bounds[1][1], res)
X, Y = np.meshgrid(x, y)
xx = np.ascontiguousarray(np.vstack([X.reshape(X.size), Y.reshape(Y.size)]).T)
model_mean = surrogate.predict(xx)
# TODO: maybe change this one below if __call__ method added to GP_acquisition
acq_value = acquisition(xx, surrogate, eval_gradient=False)
# maybe show the next max of acquisition
acq_max = xx[np.argmax(acq_value)]
fig, ax = plt.subplots(1, 2, figsize=(8, 4))
cmap = [plt.get_cmap("magma"), plt.get_cmap("viridis")]
label = ["Model mean (log-posterior)", "Acquisition function value"]
for i, Z in enumerate([model_mean, acq_value]):
ax[i].set_title(label[i])
# Boost the upper limit to avoid truncation errors.
Z = np.clip(Z, min(Z[np.isfinite(Z)]), max(Z[np.isfinite(Z)]))
levels = np.arange(min(Z) * 0.99, max(Z) * 1.01, (max(Z) - min(Z)) / 500)
Z = Z.reshape(*X.shape)
norm = cm.colors.Normalize(vmax=Z.max(), vmin=Z.min())
# # Background of the same color as the bottom of the colormap, to avoid "gaps"
# plt.gca().set_facecolor(cmap[i].colors[0])
ax[i].contourf(X, Y, Z, levels, cmap=plt.get_cmap(cmap[i], 256), norm=norm)
points = ax[i].scatter(
*surrogate.X_regress.T, edgecolors="deepskyblue", marker=r"$\bigcirc$"
)
# Plot position of next best sample
point_max = ax[i].scatter(*acq_max, marker="x", color="k")
if last_points is not None:
points_last = ax[i].scatter(
*last_points.T, edgecolors="violet", marker=r"$\bigcirc$"
)
# Bounds
ax[i].set_xlim(bounds[0][0], bounds[0][1])
ax[i].set_ylim(bounds[1][0], bounds[1][1])
# Remove ticks, for ilustrative purposes only
# ax[i].set_xticks([], minor=[])
# ax[i].set_yticks([], minor=[])
legend_labels = {points: "Training points"}
if last_points is not None:
legend_labels[points_last] = "Points added in last iteration."
legend_labels[point_max] = "Next optimal location"
fig.legend(
list(legend_labels), list(legend_labels.values()), loc="lower center", ncol=99
)
plt.subplots_adjust(left=0.1, right=0.9, bottom=0.15)
def _plot_2d_model_acquisition_finite(
surrogate, acquisition, last_points=None, res=200
):
"""
Contour plots for model prediction and acquisition function value of a 2d model.
If ``last_points`` passed, they are highlighted.
"""
if surrogate.d != 2:
warnings.warn("This plots are only possible in 2d.")
return
# TODO: option to restrict bounds to the min square containing traning samples,
# with some padding
bounds = surrogate.bounds
x = np.linspace(bounds[0][0], bounds[0][1], res)
y = np.linspace(bounds[1][0], bounds[1][1], res)
X, Y = np.meshgrid(x, y)
xx = np.ascontiguousarray(np.vstack([X.reshape(X.size), Y.reshape(Y.size)]).T)
model_mean = surrogate.predict(xx)
# TODO: maybe change this one below if __call__ method added to GP_acquisition
acq_value = acquisition(xx, surrogate, eval_gradient=False)
# maybe show the next max of acquisition
acq_max = xx[np.argmax(acq_value)]
fig, ax = plt.subplots(1, 2, figsize=(8, 4))
cmap = [plt.get_cmap("magma"), plt.get_cmap("viridis")]
label = ["Model mean (log-posterior)", "Acquisition function value"]
for i, Z in enumerate([model_mean, acq_value]):
ax[i].set_title(label[i])
# Boost the upper limit to avoid truncation errors.
Z_finite = Z[np.isfinite(Z)]
# Z_clipped = np.clip(Z_finite, min(Z[np.isfinite(Z)]), max(Z[np.isfinite(Z)]))
# Z_sort = np.sort(Z_finite)[::-1]
top_x_perc = np.sort(Z_finite)[::-1][: int(len(Z_finite) * 0.5)]
relevant_range = max(top_x_perc) - min(top_x_perc)
levels = np.linspace(
max(Z_finite) - 1.99 * relevant_range,
max(Z_finite) + 0.01 * relevant_range,
500,
)
Z[np.isfinite(Z)] = np.clip(Z_finite, min(levels), max(levels))
Z = Z.reshape(*X.shape)
norm = cm.colors.Normalize(vmax=max(levels), vmin=min(levels))
ax[i].set_facecolor("grey")
# # Background of the same color as the bottom of the colormap, to avoid "gaps"
# plt.gca().set_facecolor(cmap[i].colors[0])
ax[i].contourf(X, Y, Z, levels, cmap=plt.get_cmap(cmap[i], 256), norm=norm)
points = ax[i].scatter(
*surrogate.X_regress.T, edgecolors="deepskyblue", marker=r"$\bigcirc$"
)
# Plot position of next best sample
point_max = ax[i].scatter(*acq_max, marker="x", color="k")
if last_points is not None:
points_last = ax[i].scatter(
*last_points.T, edgecolors="violet", marker=r"$\bigcirc$"
)
# Bounds
ax[i].set_xlim(bounds[0][0], bounds[0][1])
ax[i].set_ylim(bounds[1][0], bounds[1][1])
# Remove ticks, for ilustrative purposes only
# ax[i].set_xticks([], minor=[])
# ax[i].set_yticks([], minor=[])
legend_labels = {points: "Training points"}
if last_points is not None:
legend_labels[points_last] = "Points added in last iteration."
legend_labels[point_max] = "Next optimal location"
fig.legend(
list(legend_labels), list(legend_labels.values()), loc="lower center", ncol=99
)
plt.subplots_adjust(left=0.1, right=0.9, bottom=0.15)
def _plot_2d_model_acquisition_std(surrogate, acquisition, last_points=None, res=200):
"""
Contour plots for model prediction and acquisition function value of a 2d model.
If ``last_points`` passed, they are highlighted.
"""
if surrogate.d != 2:
warnings.warn("This plots are only possible in 2d.")
return
# TODO: option to restrict bounds to the min square containing traning samples,
# with some padding
bounds = surrogate.bounds
x = np.linspace(bounds[0][0], bounds[0][1], res)
y = np.linspace(bounds[1][0], bounds[1][1], res)
X, Y = np.meshgrid(x, y)
xx = np.ascontiguousarray(np.vstack([X.reshape(X.size), Y.reshape(Y.size)]).T)
model_mean, model_std = surrogate.predict(xx, return_std=True)
# TODO: maybe change this one below if __call__ method added to GP_acquisition
acq_value = acquisition(xx, surrogate, eval_gradient=False)
# maybe show the next max of acquisition
acq_max = xx[np.argmax(acq_value)]
fig, ax = plt.subplots(1, 3, figsize=(12, 4))
cmap = [plt.get_cmap("magma"), plt.get_cmap("viridis"), plt.get_cmap("magma")]
label = [
"Model mean (log-posterior)",
"Acquisition function value",
"Model std dev.",
]
for i, Z in enumerate([model_mean, acq_value]):
ax[i].set_title(label[i])
# Boost the upper limit to avoid truncation errors.
Z_finite = Z[np.isfinite(Z)]
# Z_clipped = np.clip(Z_finite, min(Z[np.isfinite(Z)]), max(Z[np.isfinite(Z)]))
# Z_sort = np.sort(Z_finite)[::-1]
top_x_perc = np.sort(Z_finite)[::-1][: int(len(Z_finite) * 0.5)]
relevant_range = max(top_x_perc) - min(top_x_perc)
levels = np.linspace(
max(Z_finite) - 1.99 * relevant_range,
max(Z_finite) + 0.01 * relevant_range,
500,
)
Z[np.isfinite(Z)] = np.clip(Z_finite, min(levels), max(levels))
Z = Z.reshape(*X.shape)
norm = cm.colors.Normalize(vmax=max(levels), vmin=min(levels))
ax[i].set_facecolor("grey")
# # Background of the same color as the bottom of the colormap, to avoid "gaps"
# plt.gca().set_facecolor(cmap[i].colors[0])
ax[i].contourf(X, Y, Z, levels, cmap=cm.get_cmap(cmap[i], 256), norm=norm)
points = ax[i].scatter(
*surrogate.X_regress.T, edgecolors="deepskyblue", marker=r"$\bigcirc$"
)
# Plot position of next best sample
point_max = ax[i].scatter(*acq_max, marker="x", color="k")
if last_points is not None:
points_last = ax[i].scatter(
*last_points.T, edgecolors="violet", marker=r"$\bigcirc$"
)
# Bounds
ax[i].set_xlim(bounds[0][0], bounds[0][1])
ax[i].set_ylim(bounds[1][0], bounds[1][1])
# Remove ticks, for ilustrative purposes only
# ax[i].set_xticks([], minor=[])
# ax[i].set_yticks([], minor=[])
ax[2].set_title(label[2])
Z = model_std
Z_finite = Z[np.isfinite(model_mean)]
Z[~np.isfinite(model_mean)] = -np.inf
minz = min(Z_finite)
zrange = max(Z_finite) - minz
levels = np.linspace(minz, minz + (zrange if zrange > 0 else 0.00001), 500)
# Z[np.isfinite(model_mean)] = np.clip(Z_finite, min(levels), max(levels))
Z = Z.reshape(*X.shape)
norm = cm.colors.Normalize(vmax=max(levels), vmin=min(levels))
ax[2].set_facecolor("grey")
ax[2].contourf(X, Y, Z, levels, cmap=plt.get_cmap(cmap[2], 256), norm=norm)
points = ax[2].scatter(
*surrogate.X_regress.T, edgecolors="deepskyblue", marker=r"$\bigcirc$"
)
# Plot position of next best sample
point_max = ax[2].scatter(*acq_max, marker="x", color="k")
if last_points is not None:
points_last = ax[2].scatter(
*last_points.T, edgecolors="violet", marker=r"$\bigcirc$"
)
# Bounds
ax[2].set_xlim(bounds[0][0], bounds[0][1])
ax[2].set_ylim(bounds[1][0], bounds[1][1])
legend_labels = {points: "Training points"}
if last_points is not None:
legend_labels[points_last] = "Points added in last iteration."
legend_labels[point_max] = "Next optimal location"
fig.legend(
list(legend_labels), list(legend_labels.values()), loc="lower center", ncol=99
)
plt.subplots_adjust(left=0.1, right=0.9, bottom=0.15)
