import copy
import logging
from typing import Any, Callable
import array_api_compat.numpy as np
import array_api_extra as xpx
from orng import ArrayRNG
from ...flows.base import Flow
from ...history import SMCHistory
from ...samples import SMCSamples
from ...utils import (
asarray,
determine_backend_name,
effective_sample_size,
to_numpy,
track_calls,
)
from ..mcmc import MCMCSampler
[docs]
logger = logging.getLogger(__name__)
[docs]
DEFAULT_BETA_TOLERANCE = 1e-8
[docs]
class BetaScheduleError(RuntimeError):
pass
[docs]
class SMCSampler(MCMCSampler):
"""Base class for Sequential Monte Carlo samplers.
Parameters
----------
log_likelihood : Callable
The log likelihood function.
log_prior : Callable
The log prior function.
dims : int
The number of dimensions.
prior_flow : Flow
The prior flow.
xp : Callable
The array API backend.
dtype : Any | str | None, optional
The data type for the samples, by default None.
parameters : list[str] | None, optional
The parameter names, by default None.
rng : np.random.Generator | ArrayRNG | None, optional
The random number generator, by default None.
preconditioning_transform : Callable | None, optional
The preconditioning transform, by default None.
"""
def __init__(
self,
log_likelihood: Callable,
log_prior: Callable,
dims: int,
prior_flow: Flow,
xp: Callable,
dtype: Any | str | None = None,
parameters: list[str] | None = None,
rng: np.random.Generator | ArrayRNG | None = None,
preconditioning_transform: Callable | None = None,
):
super().__init__(
log_likelihood=log_likelihood,
log_prior=log_prior,
dims=dims,
prior_flow=prior_flow,
xp=xp,
dtype=dtype,
parameters=parameters,
preconditioning_transform=preconditioning_transform,
)
[docs]
self.rng = rng or ArrayRNG(determine_backend_name(xp=xp))
self._adaptive_target_efficiency = False
@property
[docs]
def target_efficiency(self):
return self._target_efficiency
@target_efficiency.setter
def target_efficiency(self, value: float | tuple):
"""Set the target efficiency.
Parameters
----------
value : float or tuple
If a float, the target efficiency to use for all iterations.
If a tuple of two floats, the target efficiency will adapt from
the first value to the second value over the course of the SMC
iterations. See `target_efficiency_rate` for details.
"""
if isinstance(value, float):
if not (0 < value < 1):
raise ValueError("target_efficiency must be in (0, 1)")
self._target_efficiency = value
self._adaptive_target_efficiency = False
elif len(value) != 2:
raise ValueError(
"target_efficiency must be a float or tuple of two floats"
)
else:
value = tuple(map(float, value))
if not (0 < value[0] < value[1] < 1):
raise ValueError(
"target_efficiency tuple must be in (0, 1) and increasing"
)
self._target_efficiency = value
self._adaptive_target_efficiency = True
[docs]
def current_target_efficiency(self, beta: float) -> float:
"""Get the current target efficiency based on beta."""
if self._adaptive_target_efficiency:
return self._target_efficiency[0] + (
self._target_efficiency[1] - self._target_efficiency[0]
) * (beta**self.target_efficiency_rate)
else:
return self._target_efficiency
[docs]
def determine_beta(
self,
samples: SMCSamples,
beta: float,
beta_step: float,
min_beta_step: float,
max_beta_step: float = 1.0,
beta_tolerance: float = DEFAULT_BETA_TOLERANCE,
) -> tuple[float, float]:
"""Determine the next beta value.
Parameters
----------
samples : SMCSamples
The current samples.
beta : float
The current beta value.
beta_step : float
The fixed beta step size if not adaptive.
min_beta_step : float
The minimum beta step size.
max_beta_step : float
The maximum beta step size.
beta_tolerance : float
Tolerance when checking for beta convergence.
Returns
-------
beta : float
The new beta value.
min_beta_step : float
The new minimum beta step size if adaptive_min_beta_step is True.
Raises
------
BetaScheduleError
If adaptive beta is enabled and the determined beta does not increase
from the previous beta.
"""
if not self.adaptive:
beta += beta_step
if beta >= 1.0:
beta = 1.0
else:
beta_prev = beta
beta_min = beta_prev
beta_max = 1.0
eff_beta_max = effective_sample_size(
samples.log_weights(beta_max)
) / len(samples)
current_eff = self.current_target_efficiency(beta_prev)
if eff_beta_max >= current_eff:
beta_min = 1.0
target_eff = current_eff
while beta_max - beta_min > beta_tolerance:
beta_try = 0.5 * (beta_max + beta_min)
eff = effective_sample_size(
samples.log_weights(beta_try)
) / len(samples)
if eff >= target_eff:
beta_min = beta_try
else:
beta_max = beta_try
beta_star = beta_min
if beta_star <= beta_prev + beta_tolerance and beta_prev < 1.0:
logger.warning(
"Adaptive beta search could not find a beta above %.6g "
"that satisfies the target efficiency %.3f within "
"tolerance %.1e; beta may remain unchanged. "
"Consider decreasing beta_tolerance or target_efficiency.",
beta_prev,
target_eff,
beta_tolerance,
)
if self.adaptive_min_beta_step:
min_beta_step = (
min_beta_step * (1 - beta_prev) / (1 - beta_star)
)
beta = max(beta_star, beta_prev + min_beta_step)
beta = min(beta, beta_prev + max_beta_step, 1.0)
if beta == beta_prev:
raise BetaScheduleError(
f"Beta did not increase from previous value {beta:.6g}. "
"Adaptive beta search may have failed to find a suitable beta. "
f"Consider adjusting beta_tolerance ({beta_tolerance}), "
f"min_beta_step ({min_beta_step}) or "
f"target_efficiency ({target_eff}) "
"(values may be adaptive)."
)
return beta, min_beta_step
@track_calls
[docs]
def sample(
self,
n_samples: int,
n_steps: int | None = None,
adaptive: bool = True,
min_beta_step: float | None = None,
max_beta_step: float | None = None,
max_n_steps: int | None = None,
target_efficiency: float = 0.5,
target_efficiency_rate: float = 1.0,
n_final_samples: int | None = None,
checkpoint_callback: Callable[[dict], None] | None = None,
checkpoint_every: int | None = None,
checkpoint_file_path: str | None = None,
resume_from: str | bytes | dict | None = None,
store_sample_history: bool = True,
beta_tolerance: float = DEFAULT_BETA_TOLERANCE,
) -> SMCSamples:
"""Sample using the SMC sampler.
Parameters
----------
n_samples : int
The number of samples (particles) to use in the SMC sampler.
n_steps : int, optional
The number of SMC iterations to perform. Must be specified if
:code:`adaptive=False`. Default is None.
adaptive : bool, optional
Whether to adaptively determine the beta schedule. Default is True.
min_beta_step : float, optional
The minimum beta step size when using adaptive beta. Default is None,
which means no minimum step size.
max_beta_step : float, optional
The maximum beta step size when using adaptive beta. Default is None,
which means no maximum step size.
max_n_steps : int, optional
The maximum number of SMC iterations to perform when using adaptive
beta. Default is None, which means no maximum.
target_efficiency : float or tuple, optional
The target sample efficiency (ESS / n_samples) to aim for at each
SMC iteration. Can be a single float in (0, 1) or a tuple of two
floats specifying a range to adapt between. Default is 0.5.
target_efficiency_rate : float, optional
When using a tuple for target_efficiency, this controls the rate at
which the target efficiency adapts from the first value to the
second value as beta increases. Default is 1.0 (linear adaptation).
n_final_samples : int, optional
If specified, the number of final samples to produce after the SMC
iterations. If not specified, the number of final samples will be
the same as n_samples. Default is None.
checkpoint_callback : callable, optional
A callback function to call with a checkpoint dictionary at regular
intervals during sampling. Default is None (no checkpointing).
checkpoint_every : int, optional
The number of iterations between checkpoints when using checkpoint
callback.
Default is None, which means no regular checkpointing.
checkpoint_file_path : str, optional
If using checkpoint_callback, this can be used to specify a file
path to save checkpoints to. Default is None.
resume_from : str, bytes, or dict, optional
If specified, this can be used to resume sampling from a previous
checkpoint. Can be a file path, bytes object, or checkpoint
dictionary. Default is None (start from scratch).
store_sample_history : bool, optional
Whether to store the history of samples at each iteration in
:code:`self.history.sample_history`. Default is True.
beta_tolerance : float, optional
Tolerance for determining convergence of beta when using adaptive
beta. Default is given by :code:`DEFAULT_BETA_TOLERANCE`.
Returns
-------
final_samples : SMCSamples
The final samples after running the SMC sampler, with log evidence
and log evidence error estimates.
Raises
------
ValueError
If both n_steps is None and adaptive is False, or if
target_efficiency is not in (0, 1) when a float, or if
target_efficiency tuple is not valid, or if log probabilities
contain NaN values.
"""
resumed = resume_from is not None
if resumed:
resume_from_printable = (
resume_from
if isinstance(resume_from, str)
else "checkpoint data"
)
logger.info(
f"Resuming SMC sampling from checkpoint: {resume_from_printable}"
)
samples, beta, iterations = self.restore_from_checkpoint(
resume_from
)
logger.info(
f"Resumed SMC sampling at iteration {iterations} with beta={beta:.4f}"
)
else:
samples = self.draw_initial_samples(n_samples)
samples = SMCSamples.from_samples(
samples, xp=self.xp, beta=0.0, dtype=self.dtype
)
beta = 0.0
iterations = 0
self.history = SMCHistory()
self.fit_preconditioning_transform(samples.x)
if store_sample_history:
self.history.sample_history.append(samples.to_numpy())
if self.xp.isnan(samples.log_q).any():
raise ValueError("Log proposal contains NaN values")
if self.xp.isnan(samples.log_prior).any():
raise ValueError("Log prior contains NaN values")
if self.xp.isnan(samples.log_likelihood).any():
raise ValueError("Log likelihood contains NaN values")
logger.debug(f"Initial sample summary: {samples}")
# Remove the n_final_steps from sampler_kwargs if present
self.sampler_kwargs = self.sampler_kwargs or {}
n_final_steps = self.sampler_kwargs.pop("n_final_steps", None)
self.target_efficiency = target_efficiency
self.target_efficiency_rate = target_efficiency_rate
if n_steps is not None:
beta_step = 1 / n_steps
elif not adaptive:
raise ValueError("Either n_steps or adaptive=True must be set")
else:
beta_step = np.nan
self.adaptive = adaptive
if min_beta_step is None:
if max_n_steps is None:
min_beta_step = 0.0
self.adaptive_min_beta_step = False
else:
min_beta_step = 1 / max_n_steps
self.adaptive_min_beta_step = True
else:
self.adaptive_min_beta_step = False
if max_beta_step is not None:
if max_beta_step <= 0 or max_beta_step >= 1:
raise ValueError("max_beta_step must be in (0, 1)")
self.max_beta_step = max_beta_step
else:
self.max_beta_step = 1.0
iterations = iterations or 0
if checkpoint_callback is None and checkpoint_every is not None:
checkpoint_callback = self.default_file_checkpoint_callback(
checkpoint_file_path
)
if checkpoint_callback is not None and checkpoint_every is None:
checkpoint_every = 1
run_smc_loop = True
if resumed:
last_beta = self.history.beta[-1] if self.history.beta else beta
if last_beta >= 1.0:
run_smc_loop = False
logger.info(
f"Checkpoint beta {last_beta:.4f} indicates SMC loop already completed, skipping to final mutation steps if needed"
)
def maybe_checkpoint(force: bool = False):
if checkpoint_callback is None:
return
should_checkpoint = force or (
checkpoint_every is not None
and checkpoint_every > 0
and iterations % checkpoint_every == 0
)
if not should_checkpoint:
return
state = self.build_checkpoint_state(samples, iterations, beta)
checkpoint_callback(state)
if run_smc_loop:
while True:
iterations += 1
beta, min_beta_step = self.determine_beta(
samples,
beta,
beta_step,
min_beta_step,
max_beta_step=self.max_beta_step,
beta_tolerance=beta_tolerance,
)
self.history.eff_target.append(
float(self.current_target_efficiency(beta))
)
logger.info(f"it {iterations} - beta: {beta}")
self.history.beta.append(float(beta))
ess = effective_sample_size(samples.log_weights(beta))
eff = ess / len(samples)
if eff < 0.1:
logger.warning(
f"it {iterations} - Low sample efficiency: {eff:.2f}"
)
self.history.ess.append(float(ess))
logger.info(
f"it {iterations} - ESS: {ess:.1f} ({eff:.2f} efficiency)"
)
self.history.ess_target.append(
float(effective_sample_size(samples.log_weights(1.0)))
)
log_evidence_ratio = samples.log_evidence_ratio(beta)
log_evidence_ratio_var = samples.log_evidence_ratio_variance(
beta
)
self.history.log_norm_ratio.append(float(log_evidence_ratio))
self.history.log_norm_ratio_var.append(
float(log_evidence_ratio_var)
)
logger.info(
f"it {iterations} - Log evidence ratio: {log_evidence_ratio:.2f} +/- {np.sqrt(log_evidence_ratio_var):.2f}"
)
samples = samples.resample(beta, rng=self.rng)
samples = self.mutate(samples, beta)
if store_sample_history:
self.history.sample_history.append(samples.to_numpy())
maybe_checkpoint()
if beta == 1.0 or (
max_n_steps is not None and iterations >= max_n_steps
):
break
# If n_final_samples is specified and differs, perform additional mutation steps
if n_final_samples is not None and len(samples.x) != n_final_samples:
logger.info(f"Generating {n_final_samples} final samples")
if not samples.xp.isfinite(samples.log_likelihood).all():
logger.warning(
"Final samples contain non-finite log likelihood values"
)
if not samples.xp.isfinite(samples.log_prior).all():
logger.warning(
"Final samples contain non-finite log prior values"
)
if not samples.xp.isfinite(samples.log_q).all():
logger.warning(
"Final samples contain non-finite log proposal values"
)
final_samples = samples.resample(
1.0, n_samples=n_final_samples, rng=self.rng
)
samples = self.mutate(final_samples, 1.0, n_steps=n_final_steps)
samples.log_evidence = samples.xp.sum(
asarray(self.history.log_norm_ratio, self.xp)
)
samples.log_evidence_error = samples.xp.sqrt(
samples.xp.sum(asarray(self.history.log_norm_ratio_var, self.xp))
)
maybe_checkpoint(force=True)
final_samples = samples.to_standard_samples()
logger.info(
f"Log evidence: {final_samples.log_evidence:.2f} +/- {final_samples.log_evidence_error:.2f}"
)
return final_samples
[docs]
def config_dict(self, include_sample_calls: str | bool = "last") -> dict:
dictionary = super().config_dict(include_sample_calls)
# Remove resume_from from the config dict if present, since it may not
# be serializable and is not needed to reconstruct the sampler
if "sample_calls" in dictionary:
if include_sample_calls == "last":
dictionary["sample_calls"]["kwargs"].pop("resume_from", None)
else:
for call_id in dictionary["sample_calls"].keys():
dictionary["sample_calls"][call_id]["kwargs"].pop(
"resume_from", None
)
return dictionary
[docs]
def mutate(self, particles):
raise NotImplementedError
[docs]
def log_prob(self, z, beta=None):
x, log_abs_det_jacobian = self.preconditioning_transform.inverse(z)
samples = SMCSamples(x, xp=self.xp, beta=beta, dtype=self.dtype)
log_q = self.prior_flow.log_prob(samples.x)
samples.log_q = samples.array_to_namespace(log_q)
samples.log_prior = self.log_prior(samples)
samples.log_likelihood = self.log_likelihood(samples)
log_prob = samples.log_p_t(
beta=beta
).flatten() + samples.array_to_namespace(log_abs_det_jacobian)
log_prob = xpx.at(log_prob, self.xp.isnan(log_prob)).set(-self.xp.inf)
return log_prob
[docs]
def build_checkpoint_state(
self, samples: SMCSamples, iteration: int, beta: float
) -> dict:
"""Prepare a serializable checkpoint payload for the sampler state."""
return super().build_checkpoint_state(
samples.to_numpy(),
iteration,
meta={"beta": beta},
)
def _checkpoint_extra_state(self) -> dict:
history_copy = copy.deepcopy(self.history)
rng_state = (
self.rng.bit_generator.state
if hasattr(self.rng, "bit_generator")
else None
)
return {
"history": history_copy,
"rng_state": rng_state,
"sampler_kwargs": getattr(self, "sampler_kwargs", None),
}
[docs]
def restore_from_checkpoint(
self, source: str | bytes | dict
) -> tuple[SMCSamples, float, int]:
samples, state = super().restore_from_checkpoint(source)
meta = state.get("meta", {}) if isinstance(state, dict) else {}
beta = None
if isinstance(meta, dict):
beta = meta.get("beta", None)
if beta is None:
beta = state.get("beta", 0.0)
iteration = state.get("iteration", 0)
self.history = state.get("history", SMCHistory())
rng_state = state.get("rng_state")
if rng_state is not None and hasattr(self.rng, "bit_generator"):
self.rng.bit_generator.state = rng_state
samples = SMCSamples.from_samples(
samples, xp=self.xp, beta=beta, dtype=self.dtype
)
return samples, beta, iteration
[docs]
class NumpySMCSampler(SMCSampler):
"""SMCSampler that maps samples and log probabilities to NumPy arrays for
compatibility with numpy-only samplers"""
def __init__(
self,
log_likelihood,
log_prior,
dims,
prior_flow,
xp,
dtype=None,
parameters=None,
preconditioning_transform=None,
):
if preconditioning_transform is not None:
preconditioning_transform = preconditioning_transform.new_instance(
xp=np
)
super().__init__(
log_likelihood,
log_prior,
dims,
prior_flow=prior_flow,
xp=xp,
dtype=dtype,
parameters=parameters,
preconditioning_transform=preconditioning_transform,
)
[docs]
def log_prob(self, z, beta=None):
log_prob = super().log_prob(z, beta=beta)
return to_numpy(log_prob)