from __future__ import annotations
import copy
from dataclasses import dataclass, field
import matplotlib.pyplot as plt
from matplotlib.figure import Figure
from .samples import SMCSamples
from .utils import load_from_h5_file, recursively_save_to_h5_file
@dataclass
[docs]
class History:
"""Base class for storing history of a sampler."""
[docs]
def save(self, h5_file, path="history"):
"""Save the history to an HDF5 file."""
dictionary = copy.deepcopy(self.__dict__)
recursively_save_to_h5_file(h5_file, path, dictionary)
@classmethod
[docs]
def load(cls, h5_file, path="history"):
"""Load the history from an HDF5 file.
Parameters
----------
h5_file : h5py.File
The open HDF5 file to load from.
path : str, optional
The path within the HDF5 file to load the history from.
Default is "history".
"""
dictionary = load_from_h5_file(h5_file, path)
# Dataclass may have fields not present in the init signature, so we
# filter the loaded dictionary to only include fields that are defined
# in the dataclass
field_names = {
field.name for field in cls.__dataclass_fields__.values()
}
filtered_dict = {
k: v for k, v in dictionary.items() if k in field_names
}
instance = cls(**filtered_dict)
for k, v in dictionary.items():
if k not in field_names:
setattr(instance, k, v)
return instance
@dataclass
[docs]
class FlowHistory(History):
[docs]
training_loss: list[float] = field(default_factory=list)
[docs]
validation_loss: list[float] = field(default_factory=list)
[docs]
def plot_loss(self) -> Figure:
"""Plot the training and validation loss."""
fig = plt.figure()
plt.plot(self.training_loss, label="Training loss")
plt.plot(self.validation_loss, label="Validation loss")
plt.legend()
plt.xlabel("Epoch")
plt.ylabel("Loss")
return fig
[docs]
def save(self, h5_file, path="flow_history"):
"""Save the history to an HDF5 file."""
super().save(h5_file, path=path)
@dataclass
[docs]
class SMCHistory(History):
[docs]
log_norm_ratio: list[float] = field(default_factory=list)
[docs]
log_norm_ratio_var: list[float] = field(default_factory=list)
[docs]
beta: list[float] = field(default_factory=list)
[docs]
ess: list[float] = field(default_factory=list)
[docs]
ess_target: list[float] = field(default_factory=list)
[docs]
eff_target: list[float] = field(default_factory=list)
[docs]
mcmc_autocorr: list[float] = field(default_factory=list)
[docs]
mcmc_acceptance: list[float] = field(default_factory=list)
[docs]
sample_history: list[SMCSamples] = field(default_factory=list)
[docs]
def save(self, h5_file, path="smc_history"):
"""Save the history to an HDF5 file.
The sample history is saved as a separate group under the main history
group, with one subgroup per iteration. The number of iterations is
stored in the main history group to allow for loading the sample history
correctly.
Parameters
----------
h5_file : h5py.File
The open HDF5 file to save to.
path : str, optional
The path within the HDF5 file to save the history. Default is
"smc_history".
"""
# Exclude sample_history from the main dictionary since it is saved separately
exclude = {"sample_history"}
dictionary = {
k: copy.deepcopy(v)
for k, v in self.__dict__.items()
if k not in exclude
}
dictionary["__len_sample_history"] = len(self.sample_history)
recursively_save_to_h5_file(h5_file, path, dictionary)
for i, samples in enumerate(self.sample_history):
samples.save(h5_file, path=f"{path}__sample_history/{i}")
@classmethod
[docs]
def load(cls, h5_file, path="smc_history"):
"""Load the history from an HDF5 file.
Parameters
----------
h5_file : h5py.File
The open HDF5 file to load from.
path : str, optional
The path within the HDF5 file to load the history from. Default is
"smc_history".
Returns
-------
SMCHistory
The loaded history instance.
"""
dictionary = load_from_h5_file(h5_file, path)
n_samples = int(dictionary.pop("__len_sample_history", 0))
dictionary["sample_history"] = [
SMCSamples.load(h5_file, path=f"{path}__sample_history/{i}")
for i in range(n_samples)
]
# Dataclass may have fields not present in the init signature, so we
# filter the loaded dictionary to only include fields that are defined
# in the dataclass
field_names = {
field.name for field in cls.__dataclass_fields__.values()
}
filtered_dict = {
k: v for k, v in dictionary.items() if k in field_names
}
instance = cls(**filtered_dict)
# Set any additional attributes that were not part of the dataclass
# fields
for k, v in dictionary.items():
if k not in field_names:
setattr(instance, k, v)
return instance
[docs]
def plot_beta(self, ax=None) -> Figure | None:
if ax is None:
fig, ax = plt.subplots()
else:
fig = None
ax.plot(self.beta)
ax.set_xlabel("Iteration")
ax.set_ylabel(r"$\beta$")
return fig
[docs]
def plot_log_norm_ratio(self, ax=None) -> Figure | None:
if ax is None:
fig, ax = plt.subplots()
else:
fig = None
ax.plot(self.log_norm_ratio)
ax.set_xlabel("Iteration")
ax.set_ylabel("Log evidence ratio")
return fig
[docs]
def plot_ess(self, ax=None) -> Figure | None:
if ax is None:
fig, ax = plt.subplots()
else:
fig = None
ax.plot(self.ess)
ax.set_xlabel("Iteration")
ax.set_ylabel("ESS")
return fig
[docs]
def plot_ess_target(self, ax=None) -> Figure | None:
if ax is None:
fig, ax = plt.subplots()
else:
fig = None
ax.plot(self.ess_target)
ax.set_xlabel("Iteration")
ax.set_ylabel("ESS target")
return fig
[docs]
def plot_eff_target(self, ax=None) -> Figure | None:
if ax is None:
fig, ax = plt.subplots()
else:
fig = None
ax.plot(self.eff_target)
ax.set_xlabel("Iteration")
ax.set_ylabel("Efficiency target")
return fig
[docs]
def plot_mcmc_acceptance(self, ax=None) -> Figure | None:
if ax is None:
fig, ax = plt.subplots()
else:
fig = None
ax.plot(self.mcmc_acceptance)
ax.set_xlabel("Iteration")
ax.set_ylabel("MCMC Acceptance")
return fig
[docs]
def plot_mcmc_autocorr(self, ax=None) -> Figure | None:
if ax is None:
fig, ax = plt.subplots()
else:
fig = None
ax.plot(self.mcmc_autocorr)
ax.set_xlabel("Iteration")
ax.set_ylabel("MCMC Autocorr")
return fig
[docs]
def plot(self, fig: Figure | None = None) -> Figure:
methods = [
self.plot_beta,
self.plot_log_norm_ratio,
self.plot_ess,
self.plot_ess_target,
self.plot_eff_target,
self.plot_mcmc_acceptance,
]
if fig is None:
fig, axs = plt.subplots(len(methods), 1, sharex=True)
else:
axs = fig.axes
for method, ax in zip(methods, axs):
method(ax)
for ax in axs[:-1]:
ax.set_xlabel("")
return fig
[docs]
def plot_sample_history(
self,
n_samples=None,
parameters=None,
ax=None,
cmap: str = "viridis",
scatter_kwargs=None,
x_axis: str = "log_p_t",
iterations: list[int] | None = None,
) -> Figure | None:
"""Plot the history of samples in the SMC sampler.
Parameters
----------
n_samples : int, optional
Number of samples to plot from each iteration. If None, plot all samples.
parameters : list of str, optional
List of parameter names to plot. If None, plot all parameters.
ax : matplotlib.axes.Axes, optional
Axes to plot on. If None, a new figure and axes will be created.
cmap : str, optional
Colormap to use for plotting the samples. Default is "viridis".
scatter_kwargs : dict, optional
Keyword arguments to pass to the scatter function. If None, no additional arguments will be passed
x_axis : str, optional
Quantity to use for the x-axis. Supported values are
:code:`"log_p_t"` and :code:`"log_likelihood"`.
Falls back to iteration index if required fields are missing.
"""
import numpy as np
if x_axis not in {"log_p_t", "log_likelihood"}:
raise ValueError(
f"Unsupported x_axis '{x_axis}'. Expected 'log_p_t' or 'log_likelihood'."
)
n_parameters = (
len(parameters)
if parameters is not None
else self.sample_history[0].dims
)
if ax is None:
fig, ax = plt.subplots(n_parameters, 1, sharex=True)
ax = np.atleast_1d(ax)
else:
fig = None
cmap = plt.get_cmap(cmap)
if iterations is None:
iterations = list(range(len(self.sample_history)))
colors = cmap(np.linspace(0, 1, len(iterations)))
has_log_pt = all(
getattr(samples, "beta", None) is not None
and samples.log_likelihood is not None
and samples.log_prior is not None
and samples.log_q is not None
for samples in self.sample_history
)
has_log_likelihood = all(
samples.log_likelihood is not None
for samples in self.sample_history
)
scatter_kwargs = scatter_kwargs or {}
default_scatter_kwargs = dict(s=10)
scatter_kwargs = {**default_scatter_kwargs, **scatter_kwargs}
for it, color in zip(iterations, colors):
samples = self.sample_history[it]
samples = samples.to_numpy()
if n_samples is not None:
samples = samples[:n_samples]
if parameters is not None:
idx = [samples.parameters.index(p) for p in parameters]
x = samples.x[:, idx]
else:
x = samples.x
for i in range(x.shape[1]):
if x_axis == "log_p_t" and has_log_pt:
x_axis_values = samples.log_p_t(samples.beta)
elif x_axis == "log_likelihood" and has_log_likelihood:
x_axis_values = samples.log_likelihood
else:
x_axis_values = it * np.ones(samples.x.shape[0])
ax[i].scatter(
x_axis_values, x[:, i], color=color, **scatter_kwargs
)
parameters = parameters or samples.parameters
if parameters is None:
parameters = [f"x_{i}" for i in range(samples.x.shape[1])]
for i, p in enumerate(parameters):
ax[i].set_ylabel(p)
if x_axis == "log_p_t" and has_log_pt:
ax[-1].set_xlabel("log p_t(beta)")
elif x_axis == "log_likelihood" and has_log_likelihood:
ax[-1].set_xlabel("log likelihood")
else:
ax[-1].set_xlabel("Iteration")
return fig
[docs]
def plot_quantile_bands(
self,
parameters: list[str] | None = None,
quantile_interval: tuple[float, float] = (0.1, 0.9),
ax=None,
line_kwargs=None,
band_kwargs=None,
) -> Figure | None:
"""Plot per-parameter quantile bands vs iteration.
Parameters
----------
parameters : list[str] | None, optional
Parameters to plot. If None, all parameters are used.
quantile_interval : tuple[float, float], optional
Lower/upper quantiles to plot as a band.
ax : matplotlib.axes.Axes | list[matplotlib.axes.Axes] | None, optional
Axes to draw on. If None, creates a new figure.
line_kwargs : dict | None, optional
Keyword arguments for the median line.
band_kwargs : dict | None, optional
Keyword arguments for the quantile band fill.
"""
import numpy as np
if not self.sample_history:
raise ValueError("No sample history available to plot.")
q_low, q_high = quantile_interval
if not (0.0 <= q_low < 0.5 and 0.5 < q_high <= 1.0 and q_low < q_high):
raise ValueError(
"quantile_interval must satisfy 0 <= low < 0.5 < high <= 1."
)
first = self.sample_history[0]
all_parameters = first.parameters or [
f"x_{i}" for i in range(first.dims)
]
if parameters is None:
parameters = all_parameters
indices = [all_parameters.index(p) for p in parameters]
n_params = len(indices)
if ax is None:
fig, axs = plt.subplots(n_params, 1, sharex=True)
axs = np.atleast_1d(axs)
else:
fig = None
axs = np.atleast_1d(ax)
if len(axs) != n_params:
raise ValueError(
"Number of axes must match number of requested parameters."
)
line_kwargs = {"color": "C0", "lw": 1.5, **(line_kwargs or {})}
band_kwargs = {"color": "C0", "alpha": 0.2, **(band_kwargs or {})}
iterations = np.arange(len(self.sample_history))
medians = np.empty((len(self.sample_history), n_params))
lowers = np.empty((len(self.sample_history), n_params))
uppers = np.empty((len(self.sample_history), n_params))
for it, samples in enumerate(self.sample_history):
x_np = samples.to_numpy().x
for j, idx in enumerate(indices):
values = x_np[:, idx]
medians[it, j] = np.quantile(values, 0.5)
lowers[it, j] = np.quantile(values, q_low)
uppers[it, j] = np.quantile(values, q_high)
for j, (axis, param) in enumerate(zip(axs, parameters)):
axis.plot(iterations, medians[:, j], **line_kwargs)
axis.fill_between(
iterations, lowers[:, j], uppers[:, j], **band_kwargs
)
axis.set_ylabel(param)
axs[-1].set_xlabel("Iteration")
return fig