from __future__ import annotations
import importlib
import logging
import math
from copy import deepcopy
from dataclasses import dataclass, field, fields
from typing import TYPE_CHECKING, Any, Callable, ClassVar
import numpy as np
from array_api_compat import (
array_namespace,
)
from array_api_compat.common._typing import Array
from array_api_extra import default_dtype
from matplotlib.figure import Figure
from .utils import (
asarray,
convert_dtype,
decode_dtype,
encode_dtype,
infer_device,
logsumexp,
recursively_save_to_h5_file,
resolve_dtype,
to_numpy,
)
[docs]
logger = logging.getLogger(__name__)
if TYPE_CHECKING:
import pandas as pd
@dataclass
[docs]
class BaseSamples:
"""Class for storing samples and corresponding weights.
If :code:`xp` is not specified, all inputs will be converted to match
the array type of :code:`x`.
"""
"""Array of samples, shape (n_samples, n_dims)."""
[docs]
log_likelihood: Array | None = None
"""Log-likelihood values for the samples."""
[docs]
log_prior: Array | None = None
"""Log-prior values for the samples."""
[docs]
log_q: Array | None = None
"""Log-probability values under the proposal distribution."""
[docs]
parameters: list[str] | None = None
"""List of parameter names."""
[docs]
dtype: Any | str | None = None
"""Data type of the samples.
If None, the default dtype for the array namespace will be used.
"""
[docs]
xp: Callable | None = None
"""
The array namespace to use for the samples.
If None, the array namespace will be inferred from the type of :code:`x`.
"""
"""Device to store the samples on.
If None, the device will be inferred from the array namespace of :code:`x`.
"""
def __post_init__(self):
if self.xp is None:
self.xp = array_namespace(self.x)
# Numpy arrays need to be on the CPU before being converted
if self.dtype is not None:
self.dtype = resolve_dtype(self.dtype, self.xp)
else:
# Fall back to default dtype for the array namespace
self.dtype = default_dtype(self.xp)
self.x = self.array_to_namespace(self.x, dtype=self.dtype)
if self.device is None:
self.device = infer_device(self.x, self.xp)
if self.log_likelihood is not None:
self.log_likelihood = self.array_to_namespace(
self.log_likelihood, dtype=self.dtype
)
if self.log_prior is not None:
self.log_prior = self.array_to_namespace(
self.log_prior, dtype=self.dtype
)
if self.log_q is not None:
self.log_q = self.array_to_namespace(self.log_q, dtype=self.dtype)
if self.parameters is None:
self.parameters = [f"x_{i}" for i in range(self.dims)]
@property
[docs]
def dims(self):
"""Number of dimensions (parameters) in the samples."""
if self.x is None:
return 0
return self.x.shape[1] if self.x.ndim > 1 else 1
[docs]
def to_numpy(self, dtype: Any | str | None = None):
logger.debug("Converting samples to numpy arrays")
import array_api_compat.numpy as np
if dtype is not None:
dtype = resolve_dtype(dtype, np)
else:
dtype = convert_dtype(self.dtype, np)
return self.__class__(
x=self.x,
parameters=self.parameters,
log_likelihood=self.log_likelihood,
log_prior=self.log_prior,
log_q=self.log_q,
xp=np,
)
[docs]
def to_namespace(self, xp, dtype: Any | str | None = None):
if dtype is None:
dtype = convert_dtype(self.dtype, xp)
else:
dtype = resolve_dtype(dtype, xp)
logger.debug("Converting samples to {} namespace", xp)
return self.__class__(
x=self.x,
parameters=self.parameters,
log_likelihood=self.log_likelihood,
log_prior=self.log_prior,
log_q=self.log_q,
xp=xp,
device=self.device,
dtype=dtype,
)
[docs]
def array_to_namespace(self, x, dtype=None):
"""Convert an array to the same namespace as the samples"""
return asarray(x, self.xp, dtype=dtype, device=self.device)
[docs]
def to_dict(self, flat: bool = True, copy: bool = True):
"""Convert the samples to a dictionary.
Parameters
----------
flat : bool
If True, the samples are stored as separate keys for each parameter.
If False, the samples are stored in a "samples" key as a dictionary
of parameter arrays.
copy : bool
If True, the arrays in the dictionary are deep-copied. If False, they
are not copied and may share memory with the original samples.
Returns
-------
dict
A dictionary representation of the samples.
"""
out = {}
for f in fields(self):
name = f.name
if name in ["x", "xp"]:
continue
value = getattr(self, name)
if value is None:
out[name] = None
else:
# This could be improved
try:
out[name] = deepcopy(value) if copy else value
except Exception:
out[name] = value
out["xp"] = self.xp
samples = dict(zip(self.parameters, self.x.T, strict=True))
if flat:
out.update(samples)
else:
out["samples"] = samples
return out
@classmethod
[docs]
def from_dict(cls, dictionary):
"""Load samples from a dictionary.
The dictionary can either be in a flat format, where the samples are
stored as separate keys for each parameter, or in a nested format, where
the samples are stored in a "samples" key as a dictionary of parameter.
"""
dictionary = dictionary.copy()
if "samples" in dictionary:
samples = dictionary.pop("samples")
parameters = dictionary.pop("parameters")
if parameters is None:
parameters = sorted(samples.keys())
x = np.stack([samples[p] for p in parameters], axis=-1)
else:
parameters = dictionary.pop("parameters")
if parameters is None:
raise ValueError(
"Parameters must be provided if samples are not nested in a 'samples' key"
)
x = np.stack([dictionary[p] for p in parameters], axis=-1)
for p in parameters:
dictionary.pop(p, None)
return cls(x=x, parameters=parameters, **dictionary)
[docs]
def to_dataframe(self, include: list[str] | None = None) -> "pd.DataFrame":
"""Convert the samples to a pandas DataFrame.
Only includes samples, log_likelihood, log_prior, and log_q by default,
since additional fields have varying shapes and may not be compatible
with a DataFrame format.
Parameters
----------
include : list[str] | None
List of fields to include in the DataFrame. If None, includes x,
log_likelihood, log_prior, and log_q. x is always included
irrespective of the value of include.
Returns
-------
pd.DataFrame
A DataFrame representation of the samples.
"""
import pandas as pd
data = {}
samples = dict(zip(self.parameters, self.x.T, strict=True))
data.update(samples)
if include is None:
include = ["log_likelihood", "log_prior", "log_q"]
for key in include:
if getattr(self, key) is not None:
data[key] = getattr(self, key)
else:
data[key] = np.full(len(self.x), np.nan)
return pd.DataFrame(data)
[docs]
def plot_corner(
self,
parameters: list[str] | None = None,
fig: Figure | None = None,
**kwargs,
):
"""Plot a corner plot of the samples.
Parameters
----------
parameters : list[str] | None
List of parameters to plot. If None, all parameters are plotted.
Figure to plot on. If None, a new figure is created.
**kwargs : dict
Additional keyword arguments to pass to corner.corner(). Kwargs
are deep-copied before use.
"""
import corner
kwargs = deepcopy(kwargs)
kwargs.setdefault("labels", self.parameters)
if parameters is not None:
indices = [self.parameters.index(p) for p in parameters]
kwargs["labels"] = parameters
x = self.x[:, indices] if self.x.ndim > 1 else self.x[indices]
else:
x = self.x
fig = corner.corner(to_numpy(x), fig=fig, **kwargs)
return fig
def __str__(self):
out = (
f"No. samples: {len(self.x)}\nNo. parameters: {self.x.shape[-1]}\n"
)
return out
def _encode_for_hdf5(self, flat=True):
"""Encode the samples for storage in an HDF5 file."""
dictionary = self.to_numpy().to_dict(flat=flat)
dictionary["dtype"] = encode_dtype(self.xp, self.dtype)
dictionary["xp"] = self.xp.__name__
return dictionary
[docs]
def save(self, h5_file, path="samples", flat=False):
"""Save the samples to an HDF5 file.
This converts the samples to numpy and then to a dictionary.
Parameters
----------
h5_file : h5py.File
The HDF5 file to save to.
path : str
The path in the HDF5 file to save to.
flat : bool
If True, save the samples as a flat dictionary.
If False, save the samples as a nested dictionary.
"""
dictionary = self._encode_for_hdf5(flat=flat)
recursively_save_to_h5_file(h5_file, path, dictionary)
@classmethod
def _decode_from_dictionary(cls, dictionary):
"""Decode the samples from a dictionary loaded from an HDF5 file."""
dictionary["xp"] = importlib.import_module(dictionary["xp"])
dictionary["dtype"] = decode_dtype(
dictionary["xp"], dictionary["dtype"]
)
return cls.from_dict(dictionary)
@classmethod
[docs]
def load(cls, h5_file, path="samples"):
"""Load the samples from an HDF5 file."""
from .utils import load_from_h5_file
dictionary = load_from_h5_file(h5_file, path)
return cls._decode_from_dictionary(dictionary)
def __len__(self):
return len(self.x)
def __getitem__(self, idx) -> BaseSamples:
return self.__class__(
x=self.x[idx],
log_likelihood=self.log_likelihood[idx]
if self.log_likelihood is not None
else None,
log_prior=self.log_prior[idx]
if self.log_prior is not None
else None,
log_q=self.log_q[idx] if self.log_q is not None else None,
parameters=self.parameters,
dtype=self.dtype,
)
def __setitem__(self, idx, value: BaseSamples):
raise NotImplementedError("Setting items is not supported")
@classmethod
[docs]
def concatenate(cls, samples: list[BaseSamples]) -> BaseSamples:
"""Concatenate multiple Samples objects."""
if not samples:
raise ValueError("No samples to concatenate")
if not all(s.parameters == samples[0].parameters for s in samples):
raise ValueError("Parameters do not match")
if not all(s.xp == samples[0].xp for s in samples):
raise ValueError("Array namespaces do not match")
if not all(s.dtype == samples[0].dtype for s in samples):
raise ValueError("Dtypes do not match")
xp = samples[0].xp
return cls(
x=xp.concatenate([s.x for s in samples], axis=0),
log_likelihood=xp.concatenate(
[s.log_likelihood for s in samples], axis=0
)
if all(s.log_likelihood is not None for s in samples)
else None,
log_prior=xp.concatenate([s.log_prior for s in samples], axis=0)
if all(s.log_prior is not None for s in samples)
else None,
log_q=xp.concatenate([s.log_q for s in samples], axis=0)
if all(s.log_q is not None for s in samples)
else None,
parameters=samples[0].parameters,
dtype=samples[0].dtype,
)
@classmethod
[docs]
def from_samples(cls, samples: BaseSamples, **kwargs) -> BaseSamples:
"""Create a Samples object from a BaseSamples object."""
xp = kwargs.pop("xp", samples.xp)
device = kwargs.pop("device", samples.device)
dtype = kwargs.pop("dtype", samples.dtype)
if dtype is not None:
dtype = resolve_dtype(dtype, xp)
else:
dtype = convert_dtype(samples.dtype, xp)
return cls(
x=samples.x,
log_likelihood=samples.log_likelihood,
log_prior=samples.log_prior,
log_q=samples.log_q,
parameters=samples.parameters,
xp=xp,
device=device,
**kwargs,
)
def __getstate__(self):
state = self.__dict__.copy()
# replace xp (callable) with module name string
if self.xp is not None:
state["xp"] = (
self.xp.__name__ if hasattr(self.xp, "__name__") else None
)
return state
def __setstate__(self, state):
# Restore xp by checking the namespace of x
state["xp"] = array_namespace(state["x"])
# device may be string; leave as-is or None
device = state.get("device")
if device is not None and "jax" in getattr(
state["xp"], "__name__", ""
):
device = None
state["device"] = device
self.__dict__.update(state)
@dataclass
[docs]
class Samples(BaseSamples):
"""Class for storing samples and corresponding weights.
If :code:`xp` is not specified, all inputs will be converted to match
the array type of :code:`x`.
"""
[docs]
log_w: Array = field(init=False)
[docs]
weights: Array = field(init=False)
[docs]
evidence: float = field(init=False)
[docs]
evidence_error: float = field(init=False)
[docs]
log_evidence: float | None = None
[docs]
log_evidence_error: float | None = None
[docs]
effective_sample_size: float = field(init=False)
def __post_init__(self):
super().__post_init__()
if all(
x is not None
for x in [self.log_likelihood, self.log_prior, self.log_q]
):
self.compute_weights()
else:
self.log_w = None
self.weights = None
self.evidence = None
self.evidence_error = None
self.effective_sample_size = None
@property
[docs]
def efficiency(self):
"""Efficiency of the weighted samples.
Defined as ESS / number of samples.
"""
if self.log_w is None:
raise RuntimeError("Samples do not contain weights!")
return self.effective_sample_size / len(self.x)
[docs]
def compute_weights(self):
"""Compute the posterior weights."""
self.log_w = self.log_likelihood + self.log_prior - self.log_q
self.log_evidence = asarray(logsumexp(self.log_w), self.xp) - math.log(
len(self.x)
)
self.weights = self.xp.exp(self.log_w)
self.evidence = self.xp.exp(self.log_evidence)
n = len(self.x)
self.evidence_error = self.xp.sqrt(
self.xp.sum((self.weights - self.evidence) ** 2) / (n * (n - 1))
)
self.log_evidence_error = self.xp.abs(
self.evidence_error / self.evidence
)
log_w = self.log_w - self.xp.max(self.log_w)
self.effective_sample_size = self.xp.exp(
asarray(logsumexp(log_w) * 2 - logsumexp(log_w * 2), self.xp)
)
@property
[docs]
def scaled_weights(self):
return self.xp.exp(self.log_w - self.xp.max(self.log_w))
[docs]
def rejection_sample(self, rng=None):
if rng is None:
rng = np.random.default_rng()
log_u = asarray(
np.log(rng.uniform(size=len(self.x))), self.xp, device=self.device
)
log_w = self.log_w - self.xp.max(self.log_w)
accept = log_w > log_u
return self.__class__(
x=self.x[accept],
log_likelihood=self.log_likelihood[accept],
log_prior=self.log_prior[accept],
dtype=self.dtype,
)
[docs]
def plot_corner(self, include_weights: bool = True, **kwargs):
kwargs = deepcopy(kwargs)
if (
include_weights
and self.weights is not None
and "weights" not in kwargs
):
kwargs["weights"] = to_numpy(self.scaled_weights)
return super().plot_corner(**kwargs)
def __str__(self):
out = super().__str__()
if self.log_evidence is not None:
out += f"Log evidence: {self.log_evidence:.2f} +/- {self.log_evidence_error:.2f}\n"
if self.log_w is not None:
out += (
f"Effective sample size: {self.effective_sample_size:.1f}\n"
f"Efficiency: {self.efficiency:.2f}\n"
)
return out
[docs]
def to_namespace(self, xp):
return self.__class__(
x=asarray(self.x, xp, dtype=self.dtype),
parameters=self.parameters,
log_likelihood=asarray(self.log_likelihood, xp, dtype=self.dtype)
if self.log_likelihood is not None
else None,
log_prior=asarray(self.log_prior, xp, dtype=self.dtype)
if self.log_prior is not None
else None,
log_q=asarray(self.log_q, xp, dtype=self.dtype)
if self.log_q is not None
else None,
log_evidence=asarray(self.log_evidence, xp, dtype=self.dtype)
if self.log_evidence is not None
else None,
log_evidence_error=asarray(
self.log_evidence_error, xp, dtype=self.dtype
)
if self.log_evidence_error is not None
else None,
)
[docs]
def to_numpy(self):
return self.__class__(
x=to_numpy(self.x),
parameters=self.parameters,
log_likelihood=to_numpy(self.log_likelihood)
if self.log_likelihood is not None
else None,
log_prior=to_numpy(self.log_prior)
if self.log_prior is not None
else None,
log_q=to_numpy(self.log_q) if self.log_q is not None else None,
log_evidence=self.log_evidence
if self.log_evidence is not None
else None,
log_evidence_error=self.log_evidence_error
if self.log_evidence_error is not None
else None,
)
[docs]
def to_dataframe(self, include: list[str] | None = None) -> "pd.DataFrame":
"""Convert the samples to a pandas DataFrame.
By default, includes log_likelihood, log_prior, log_q, and log_w.
See parent class for more details.
Parameters
----------
include : list[str] | None
List of fields to include in the DataFrame. If None, includes
log_likelihood, log_prior, log_q, and log_w.
Returns
-------
pd.DataFrame
A DataFrame representation of the samples.
"""
if include is None:
include = ["log_likelihood", "log_prior", "log_q", "log_w"]
return super().to_dataframe(include)
def __getitem__(self, idx):
sliced = super().__getitem__(idx)
sliced.log_evidence = self.log_evidence
sliced.log_evidence_error = self.log_evidence_error
if self.log_w is not None:
sliced.log_w = self.array_to_namespace(self.log_w[idx])
if self.weights is not None:
sliced.weights = self.array_to_namespace(self.weights[idx])
else:
sliced.weights = sliced.xp.exp(sliced.log_w)
log_w = sliced.log_w - sliced.xp.max(sliced.log_w)
sliced.effective_sample_size = sliced.xp.exp(
asarray(logsumexp(log_w) * 2 - logsumexp(log_w * 2), sliced.xp)
)
return sliced
@dataclass
[docs]
class MCMCSamples(BaseSamples):
"""Class for storing MCMC samples and chain metadata.
Samples are stored flattened in :code:`x`, with :code:`chain_shape`
capturing the original chain layout (excluding the parameter dimension).
"""
[docs]
chain_shape: tuple[int, ...] | None = None
"""Shape of the chain excluding the parameter dimension."""
[docs]
thin: int | None = None
"""Thinning factor used to produce the chain, if any."""
[docs]
burn_in: int | None = None
"""Number of burn-in steps removed, if any."""
[docs]
autocorrelation_time: Array | None = None
"""Autocorrelation time per dimension, if available."""
[docs]
minimum_chain_ndim: ClassVar[int] = 2
"""Minimum required dimensionality for the input chain."""
def __post_init__(self):
super().__post_init__()
if self.chain_shape is None:
self.chain_shape = (len(self.x),)
else:
expected = int(np.prod(self.chain_shape))
if expected != len(self.x):
raise ValueError(
"chain_shape does not match the number of samples"
)
@classmethod
[docs]
def from_chain(
cls,
chain,
log_likelihood=None,
log_prior=None,
log_q=None,
parameters=None,
xp=None,
dtype: Any | str | None = None,
device: Any = None,
thin: int | None = None,
burn_in: int | None = None,
autocorrelation_time: float | None = None,
**extra_kwargs,
) -> "MCMCSamples":
"""Create samples from a chain array.
Parameters
----------
chain
MCMC chain with shape (..., n_dims). For a single chain, this
is typically (n_steps, n_dims). For ensemble samplers, use
(n_steps, n_walkers, n_dims).
log_likelihood, log_prior, log_q
Optional arrays matching the chain shape without the last
dimension. They will be flattened to align with :code:`x`.
parameters
Optional list of parameter names.
xp
Optional array namespace. If None, inferred from :code:`chain`.
dtype
Optional dtype to use for stored arrays.
device
Optional device to place arrays on.
thin, burn_in
Optional metadata describing how the chain was processed.
"""
if xp is None:
xp = array_namespace(chain)
chain = asarray(chain, xp)
if chain.ndim < cls.minimum_chain_ndim:
raise ValueError(
f"chain must have at least {cls.minimum_chain_ndim} dimensions"
)
chain_shape = chain.shape[:-1]
dims = chain.shape[-1]
x = xp.reshape(chain, (-1, dims))
kwargs = dict(
x=x,
log_likelihood=cls._flatten_chain_values(log_likelihood, xp),
log_prior=cls._flatten_chain_values(log_prior, xp),
log_q=cls._flatten_chain_values(log_q, xp),
parameters=parameters,
xp=xp,
device=device,
dtype=dtype,
chain_shape=chain_shape,
thin=thin,
burn_in=burn_in,
autocorrelation_time=autocorrelation_time,
)
kwargs.update(extra_kwargs)
return cls(**kwargs)
@staticmethod
def _flatten_chain_values(values, xp):
if values is None:
return None
values = asarray(values, xp)
return xp.reshape(values, (-1,))
@property
[docs]
def chain(self):
"""The chain reshaped to its original layout"""
return self.xp.reshape(self.x, (*self.chain_shape, self.dims))
@property
[docs]
def n_steps(self) -> int:
"""Number of steps in the chain, excluding walkers"""
return self.chain_shape[0]
@property
[docs]
def n_chains(self) -> int:
"""Number of parallel chains (e.g. walkers), if applicable."""
return self.chain_shape[1] if len(self.chain_shape) > 1 else 1
def _reshape_like_chain(self, values):
if values is None:
return None
return self.xp.reshape(values, self.chain_shape)
def _chain_post_process_index(self, burn_in: int, thin: int):
return slice(burn_in, None, thin)
def _post_process_constructor_kwargs(self) -> dict[str, Any]:
return {}
[docs]
def post_process(self, burn_in: int = 0, thin: int = 1) -> MCMCSamples:
"""Return a new MCMCSamples object with burn-in discarded and/or thinned."""
if burn_in < 0:
raise ValueError("burn_in must be non-negative")
if thin <= 0:
raise ValueError("thin must be a positive integer")
if burn_in == 0 and thin == 1:
logger.warning(
"No burn-in or thinning applied, returning original samples"
)
return self # No processing needed
index = self._chain_post_process_index(burn_in, thin)
chain = self.chain[index]
chain_shape = (chain.shape[0],) + chain.shape[1:-1]
x = chain.reshape((-1, self.dims))
log_likelihood = (
self._reshape_like_chain(self.log_likelihood)[index].reshape(-1)
if self.log_likelihood is not None
else None
)
log_prior = (
self._reshape_like_chain(self.log_prior)[index].reshape(-1)
if self.log_prior is not None
else None
)
log_q = (
self._reshape_like_chain(self.log_q)[index].reshape(-1)
if self.log_q is not None
else None
)
return self.__class__(
x=x,
log_likelihood=log_likelihood,
log_prior=log_prior,
log_q=log_q,
parameters=self.parameters,
xp=self.xp,
device=self.device,
dtype=self.dtype,
chain_shape=chain_shape,
thin=self.thin * thin if self.thin is not None else thin,
burn_in=(self.burn_in + burn_in)
if self.burn_in is not None
else burn_in,
autocorrelation_time=self.autocorrelation_time,
**self._post_process_constructor_kwargs(),
)
def __getitem__(self, idx):
sliced = super().__getitem__(idx)
sliced.chain_shape = (len(sliced.x),)
sliced.thin = self.thin
sliced.burn_in = self.burn_in
sliced.autocorrelation_time = self.autocorrelation_time
return sliced
def _to_numpy_constructor_kwargs(self) -> dict[str, Any]:
return {}
[docs]
def to_numpy(self):
chain_shape = None
if self.chain_shape is not None:
chain_shape = tuple(int(v) for v in np.asarray(self.chain_shape))
return self.__class__(
x=to_numpy(self.x),
parameters=self.parameters,
log_likelihood=to_numpy(self.log_likelihood)
if self.log_likelihood is not None
else None,
log_prior=to_numpy(self.log_prior)
if self.log_prior is not None
else None,
log_q=to_numpy(self.log_q) if self.log_q is not None else None,
chain_shape=chain_shape,
thin=self.thin,
burn_in=self.burn_in,
autocorrelation_time=to_numpy(self.autocorrelation_time)
if self.autocorrelation_time is not None
else None,
**self._to_numpy_constructor_kwargs(),
)
@dataclass
[docs]
class PTMCMCSamples(MCMCSamples):
"""Class for storing parallel-tempered MCMC samples."""
[docs]
betas: Array | None = None
"""Inverse temperatures for the chains.
Should be a 1D array of shape (n_temps,) in decreasing order, starting at 1.
"""
[docs]
minimum_chain_ndim: ClassVar[int] = 3
"""Minimum required dimensionality for the PTMCMC input chain."""
def __post_init__(self):
super().__post_init__()
# Ensure beta are decreasing and match the number of temperatures in the chain
if self.betas is not None:
self.betas = self.array_to_namespace(self.betas, dtype=self.dtype)
if self.betas.ndim != 1:
raise ValueError("betas must be a 1D array")
if len(self.betas) != self.n_temps:
raise ValueError(
"Length of betas must match the number of temperatures in the chain"
)
if not (self.betas[0] == 1):
raise ValueError("betas must start at 1")
if not self.xp.all(self.xp.diff(self.betas) < 0):
raise ValueError("betas must be in decreasing order")
[docs]
def subsample(
self, n_samples_per_temperature: int, rng=None
) -> "PTMCMCSamples":
"""Subsample the chain to a fixed number of samples per temperature.
Samples are drawn without replacement independently for each temperature.
Parameters
----------
n_samples_per_temperature
Number of samples to draw per temperature.
rng
Optional numpy random generator. If None, a new one is created.
"""
if rng is None:
rng = np.random.default_rng()
n_per_temp = int(np.prod(self.chain_shape[1:]))
if n_samples_per_temperature > n_per_temp:
raise ValueError(
f"n_samples_per_temperature ({n_samples_per_temperature}) exceeds "
f"available samples per temperature ({n_per_temp})"
)
# Reshape to (n_temps, n_per_temp, ...) for per-temperature indexing
chain = self.xp.reshape(self.x, (self.n_temps, n_per_temp, self.dims))
ll = (
self.xp.reshape(
self._reshape_like_chain(self.log_likelihood),
(self.n_temps, n_per_temp),
)
if self.log_likelihood is not None
else None
)
lp = (
self.xp.reshape(
self._reshape_like_chain(self.log_prior),
(self.n_temps, n_per_temp),
)
if self.log_prior is not None
else None
)
lq = (
self.xp.reshape(
self._reshape_like_chain(self.log_q),
(self.n_temps, n_per_temp),
)
if self.log_q is not None
else None
)
chains_sub, ll_sub, lp_sub, lq_sub = [], [], [], []
for t in range(self.n_temps):
idx = rng.choice(
n_per_temp, size=n_samples_per_temperature, replace=False
)
chains_sub.append(chain[t][idx])
if ll is not None:
ll_sub.append(ll[t][idx])
if lp is not None:
lp_sub.append(lp[t][idx])
if lq is not None:
lq_sub.append(lq[t][idx])
return self.__class__.from_chain(
chain=self.xp.stack(chains_sub, axis=0),
betas=self.betas,
log_likelihood=self.xp.stack(ll_sub, axis=0) if ll_sub else None,
log_prior=self.xp.stack(lp_sub, axis=0) if lp_sub else None,
log_q=self.xp.stack(lq_sub, axis=0) if lq_sub else None,
parameters=self.parameters,
xp=self.xp,
dtype=self.dtype,
device=self.device,
thin=self.thin,
burn_in=self.burn_in,
autocorrelation_time=self.autocorrelation_time,
)
@classmethod
[docs]
def from_chain(
cls,
chain,
betas=None,
log_likelihood=None,
log_prior=None,
log_q=None,
parameters=None,
xp=None,
dtype: Any | str | None = None,
device: Any = None,
thin: int | None = None,
burn_in: int | None = None,
autocorrelation_time: float | None = None,
) -> "PTMCMCSamples":
"""Create samples from a parallel-tempered MCMC chain.
Parameters
----------
chain
PTMCMC chain with shape (n_temps, ..., n_dims). A common layout
is (n_temps, n_steps, n_walkers, n_dims).
betas
Optional inverse temperatures with shape (n_temps,).
log_likelihood, log_prior, log_q
Optional arrays matching the chain shape without the last
dimension. They will be flattened to align with :code:`x`.
parameters
Optional list of parameter names.
xp
Optional array namespace. If None, inferred from :code:`chain`.
dtype
Optional dtype to use for stored arrays.
device
Optional device to place arrays on.
thin, burn_in
Optional metadata describing how the chain was processed.
"""
return super().from_chain(
chain=chain,
log_likelihood=log_likelihood,
log_prior=log_prior,
log_q=log_q,
parameters=parameters,
xp=xp,
dtype=dtype,
device=device,
betas=betas,
thin=thin,
burn_in=burn_in,
autocorrelation_time=autocorrelation_time,
)
@property
[docs]
def n_temps(self) -> int:
"""Number of temperatures in the parallel-tempered chain."""
return self.chain_shape[0]
def _chain_post_process_index(self, burn_in: int, thin: int):
return (slice(None), slice(burn_in, None, thin))
def _post_process_constructor_kwargs(self) -> dict[str, Any]:
return {"betas": self.betas}
[docs]
def at_temperature(self, index: int) -> MCMCSamples:
chain = self.chain[index]
log_likelihood = self._reshape_like_chain(self.log_likelihood)
log_prior = self._reshape_like_chain(self.log_prior)
log_q = self._reshape_like_chain(self.log_q)
return MCMCSamples.from_chain(
chain=chain,
log_likelihood=None
if log_likelihood is None
else log_likelihood[index],
log_prior=None if log_prior is None else log_prior[index],
log_q=None if log_q is None else log_q[index],
parameters=self.parameters,
xp=self.xp,
dtype=self.dtype,
device=self.device,
thin=self.thin,
burn_in=self.burn_in,
autocorrelation_time=self.autocorrelation_time[index]
if self.autocorrelation_time is not None
else None,
)
[docs]
def cold_chain(self) -> MCMCSamples:
return self.at_temperature(0)
def _to_numpy_constructor_kwargs(self) -> dict[str, Any]:
return {
"betas": to_numpy(self.betas) if self.betas is not None else None
}
[docs]
def log_evidence_thermodynamic_integration(
self, burn_in_fraction: float | None = 0.1, method: str = "variance"
) -> tuple[float, float]:
"""Compute the log evidence using thermodynamic integration.
Notes
-----
By default, follows the implementation outlined in Section 2.1.3 of
Annis et al. (2019) [1]_. If method="variance", the error is estimated
using the variance of the TI estimate across samples as in Eq. (37).
If method="coarse", the error is estimated by comparing to a coarser
integration using every other temperature, as in the original ptemcee
implementation.
.. [1] Annis, J., et al. https://doi.org/10.1016/j.jmp.2019.01.005
Parameters
----------
burn_in_fraction
Fraction of initial samples to discard as burn-in. If None, no burn-in is discarded.
This is applied independently to each temperature chain before integration.
method
Method for estimating the uncertainty in the log evidence.
Options are "variance" or "coarse". See Notes for details.
Returns
-------
log_evidence
The estimated log evidence from thermodynamic integration.
log_evidence_error
An estimate of the uncertainty in the log evidence.
"""
if self.betas is None:
raise ValueError("Betas must be provided to compute evidence")
logl_chain = self._reshape_like_chain(self.log_likelihood)
istart = (
int(logl_chain.shape[1] * burn_in_fraction)
if burn_in_fraction is not None
else 0
)
# Discard burn-in and flatten chain dimensions to a per-temperature sample axis.
logl_chain = logl_chain[:, istart:]
logl_chain = logl_chain.reshape(logl_chain.shape[0], -1)
if logl_chain.shape[1] == 0:
raise ValueError(
"No samples available after burn-in for TI evidence"
)
# Integrate from low to high temperature as in Eq. (35).
order = np.argsort(self.betas)
betas = self.betas[order]
logls = logl_chain[order]
mean_logls = np.mean(logls, axis=1)
log_evidence = np.trapezoid(mean_logls, betas)
if method == "variance":
# Eq. (36): TI_i for each aligned sample index across temperatures.
ti_per_sample = np.trapezoid(logls, betas, axis=0)
# Eq. (37): Var(mu_TI) = Var(TI) / n.
n = ti_per_sample.shape[0]
var_mu_ti = np.var(ti_per_sample) / n
log_evidence_error = math.sqrt(float(var_mu_ti))
elif method == "coarse":
# Alternative error estimate by comparing to a coarser integration using every other temperature.
# Copied from the original implementation in ptemcee
betas = betas[::-1]
logls = mean_logls[::-1]
betas0 = np.copy(betas)
if betas[-1] != 0:
logger.warning(
"Hottest chain is not at beta=0, duplicating hottest chain with beta=0 for error estimation"
)
betas = np.concatenate((betas0, [0]))
betas2 = np.concatenate((betas0[::2], [0]))
# Duplicate mean log-likelihood of hottest chain as a best guess for beta = 0.
logls2 = np.concatenate((logls[::2], [logls[-1]]))
logls = np.concatenate((logls, [logls[-1]]))
else:
betas2 = np.concatenate((betas0[:-1:2], [0]))
logls2 = np.concatenate((logls[:-1:2], [logls[-1]]))
log_evidence_2 = -np.trapezoid(logls2, betas2)
log_evidence_error = abs(log_evidence - log_evidence_2)
else:
raise ValueError(
f"Invalid method for log evidence error estimation: {method}"
)
return float(log_evidence), float(log_evidence_error)
[docs]
def log_evidence_stepping_stone(
self, burn_in_fraction: float | None = 0.1
) -> float:
"""Compute the log evidence using the stepping-stone estimator.
Notes
-----
Follows the implementation outlined in Section 2.2.3 of Annis et al. (2019) [1]_
.. [1] Annis, J., et al. https://doi.org/10.1016/j.jmp.2019.01.005
Parameters
----------
burn_in_fraction
Fraction of initial samples to discard as burn-in. If None, no burn-in is discarded.
This is applied independently to each temperature chain before integration.
Returns
-------
log_evidence
The estimated log evidence from thermodynamic integration.
log_evidence_error
An estimate of the uncertainty in the log evidence.
"""
if self.betas is None:
raise ValueError("Betas must be provided to compute evidence")
if self.betas[-1] != 0:
raise ValueError(
"Stepping stone estimator requires the hottest chain to be at beta=0"
)
logl_chain = self._reshape_like_chain(self.log_likelihood)
istart = (
int(logl_chain.shape[1] * burn_in_fraction)
if burn_in_fraction is not None
else 0
)
# Discard burn-in steps
logl_chain = logl_chain[:, istart:]
# Combine the walker and step dimensions for easier indexing but
# keep the temperature dimension separate
logl_chain = logl_chain.reshape(logl_chain.shape[0], -1)
order = np.argsort(self.betas)[::-1]
betas = self.betas[order]
logls = logl_chain[order]
log_evidence = 0.0
var_log_ss = 0.0
n_samples = logls.shape[1]
if n_samples == 0:
raise ValueError(
"No samples available after burn-in for stepping-stone evidence"
)
for i in range(len(betas) - 1):
dbeta = betas[i] - betas[i + 1] # positive
# Equation (51): log r_j = log(mean(exp(dbeta * logL)))).
a = dbeta * logls[i + 1]
a_max = np.max(a)
exp_shift = np.exp(a - a_max)
mean_shift = float(np.mean(exp_shift))
log_evidence += math.log(mean_shift) + float(a_max)
# Equation (53): Var(log SS) = (1/n^2) * sum_j sum_i (exp(a_i)/r_j)^2
ratio = exp_shift / mean_shift
var_log_ss += float(np.sum(ratio**2))
var_log_ss /= n_samples**2
return float(log_evidence), math.sqrt(float(var_log_ss))
[docs]
def plot_chain(
self,
beta_index: int,
n_walkers: int | None = None,
burn_in: int = 0,
parameters: list[str] | None = None,
fig: Figure | None = None,
**kwargs,
):
import matplotlib.pyplot as plt
chain = self.chain
if parameters is not None:
if self.parameters is None:
raise ValueError(
"Cannot specify parameters to plot if samples do not have parameter names"
)
param_indices = [self.parameters.index(p) for p in parameters]
else:
param_indices = range(chain.shape[-1])
fig, axs = plt.subplots(len(param_indices), 1, sharex=True)
for count, idx in enumerate(param_indices):
axs[count].plot(chain[beta_index, :, :n_walkers, idx], **kwargs)
axs[count].axvline(burn_in, color="r", linestyle="--")
fig.suptitle(f"$\\beta = {self.betas[beta_index]}$")
return fig
def __getitem__(self, idx):
raise NotImplementedError(
"Slicing is not supported for PTMCMCSamples. Use at_temperature()"
" to extract samples at a specific temperature."
)
@dataclass
[docs]
class SMCSamples(BaseSamples):
[docs]
beta: float | None = None
"""Temperature parameter for the current samples."""
[docs]
log_evidence: float | None = None
"""Log evidence estimate for the current samples."""
[docs]
log_evidence_error: float | None = None
"""Log evidence error estimate for the current samples."""
[docs]
def log_p_t(self, beta):
log_p_T = self.log_likelihood + self.log_prior
return (1 - beta) * self.log_q + beta * log_p_T
[docs]
def unnormalized_log_weights(self, beta: float) -> Array:
return (self.beta - beta) * self.log_q + (beta - self.beta) * (
self.log_likelihood + self.log_prior
)
[docs]
def log_evidence_ratio(self, beta: float) -> float:
log_w = self.unnormalized_log_weights(beta)
return logsumexp(log_w) - math.log(len(self.x))
[docs]
def log_evidence_ratio_variance(self, beta: float) -> float:
"""Estimate the variance of the log evidence ratio using the delta method.
Defined as Var(log Z) = Var(w) / (E[w])^2 where w are the unnormalized weights.
"""
log_w = self.unnormalized_log_weights(beta)
m = self.xp.max(log_w)
u = self.xp.exp(log_w - m)
mean_w = self.xp.mean(u)
var_w = self.xp.var(u)
return (
var_w / (len(self) * (mean_w**2)) if mean_w != 0 else self.xp.nan
)
[docs]
def log_weights(self, beta: float) -> Array:
log_w = self.unnormalized_log_weights(beta)
if self.xp.isnan(log_w).any():
raise ValueError(f"Log weights contain NaN values for beta={beta}")
log_evidence_ratio = logsumexp(log_w) - math.log(len(self.x))
return log_w + log_evidence_ratio
[docs]
def resample(
self,
beta,
n_samples: int | None = None,
rng: np.random.Generator = None,
) -> "SMCSamples":
"""Resample the samples to a new inverse temperature beta.
If :code:`beta` is the same as the current beta and :code:`n_samples`
is None or equal to the current number of samples, returns the same
samples with updated beta.
"""
if rng is None:
rng = np.random.default_rng()
if n_samples is None:
n_samples = len(self.x)
if beta == self.beta:
if n_samples is None or n_samples == len(self.x):
logger.warning(
"Resampling with the same beta value, returning identical samples"
)
return self
else:
log_w = self.xp.zeros(len(self.x))
else:
log_w = self.log_weights(beta)
w = to_numpy(self.xp.exp(log_w - logsumexp(log_w)))
idx = rng.choice(len(self.x), size=n_samples, replace=True, p=w)
return self.__class__(
x=self.x[idx],
log_likelihood=self.log_likelihood[idx],
log_prior=self.log_prior[idx],
log_q=self.log_q[idx],
beta=beta,
dtype=self.dtype,
parameters=self.parameters,
)
def __str__(self):
out = super().__str__()
if self.log_evidence is not None:
out += f"Log evidence: {self.log_evidence:.2f}\n"
return out
[docs]
def to_standard_samples(self):
"""Convert the samples to standard samples."""
return Samples(
x=self.x,
log_likelihood=self.log_likelihood,
log_prior=self.log_prior,
xp=self.xp,
parameters=self.parameters,
log_evidence=self.log_evidence,
log_evidence_error=self.log_evidence_error,
)
[docs]
def to_numpy(self):
return self.__class__(
x=to_numpy(self.x),
parameters=self.parameters,
log_likelihood=to_numpy(self.log_likelihood)
if self.log_likelihood is not None
else None,
log_prior=to_numpy(self.log_prior)
if self.log_prior is not None
else None,
log_q=to_numpy(self.log_q) if self.log_q is not None else None,
beta=self.beta,
log_evidence=self.log_evidence
if self.log_evidence is not None
else None,
log_evidence_error=self.log_evidence_error
if self.log_evidence_error is not None
else None,
)
def __getitem__(self, idx):
sliced = super().__getitem__(idx)
sliced.beta = self.beta
sliced.log_evidence = self.log_evidence
sliced.log_evidence_error = self.log_evidence_error
return sliced