import logging
from typing import Callable
import numpy as np
from orng import ArrayRNG
from ..samples import MCMCSamples, Samples, to_numpy
from ..utils import AspireFile, track_calls
from .base import Sampler
[docs]
logger = logging.getLogger(__name__)
[docs]
class MCMCSampler(Sampler):
"""Base class for MCMC samplers."""
[docs]
chain_checkpoint_path = "checkpoint"
"""Path within checkpoint file to save MCMC chain checkpoints.
The default is "checkpoint".
"""
[docs]
chain_dataset_name = "mcmc_chain"
"""Name of chain entry within checkpoint file to save MCMC checkpoints."""
def __init__(
self,
log_likelihood,
log_prior,
dims,
prior_flow,
xp,
dtype=None,
parameters=None,
preconditioning_transform=None,
rng=None,
):
super().__init__(
log_likelihood,
log_prior,
dims,
prior_flow,
xp,
dtype,
parameters,
preconditioning_transform,
)
[docs]
self.rng = rng or ArrayRNG(backend=self.backend_str)
[docs]
def draw_initial_samples(self, n_samples: int) -> Samples:
"""Draw initial samples from the prior flow.
Parameters
----------
n_samples : int
The number of samples to draw.
Returns
-------
Samples
The drawn samples, with log probabilities, log prior, and log likelihood.
"""
# Flow may propose samples outside prior bounds, so we may need
# to try multiple times to get enough valid samples.
n_samples_drawn = 0
samples = None
while n_samples_drawn < n_samples:
x, log_q = self.prior_flow.sample_and_log_prob(n_samples)
if not self.prior_flow.xp.isfinite(log_q).all():
raise ValueError(
"Proposal returned non-finite log probabilities. "
"aspire assumes the proposal is a valid, normalized "
"probability distribution and should therefore only "
"return samples with finite log probabilities."
)
new_samples = Samples(
x,
xp=self.xp,
log_q=log_q,
dtype=self.dtype,
parameters=self.parameters,
)
new_samples.log_prior = new_samples.array_to_namespace(
self.log_prior(new_samples)
)
new_samples.log_likelihood = new_samples.array_to_namespace(
self.log_likelihood(new_samples)
)
valid = self.xp.isfinite(new_samples.log_prior) & self.xp.isfinite(
new_samples.log_likelihood
)
if any(~valid):
logger.debug(
f"Proposal returned {int(~valid.sum())} invalid samples "
f"with non-finite log prior or log likelihood. "
f"These samples will be discarded."
)
n_valid = int(self.xp.sum(valid))
if n_valid > 0:
if samples is None:
samples = new_samples[valid]
else:
samples = Samples.concatenate(
[samples, new_samples[valid]]
)
n_samples_drawn += n_valid
if n_samples_drawn > n_samples:
samples = samples[:n_samples]
return samples
[docs]
def log_prob(self, z):
"""Compute the log probability of the samples.
Input samples are in the transformed space.
"""
x, log_abs_det_jacobian = self.preconditioning_transform.inverse(z)
samples = Samples(x, xp=self.xp, dtype=self.dtype)
samples.log_prior = self.log_prior(samples)
samples.log_likelihood = self.log_likelihood(samples)
log_prob = (
samples.log_likelihood
+ samples.log_prior
+ samples.array_to_namespace(log_abs_det_jacobian)
)
return log_prob.flatten()
[docs]
def default_mcmc_chain_file_checkpoint_callback(
self, file_path: str | None
) -> Callable[[dict], None]:
"""Return a callback that saves MCMC checkpoints as native HDF5 samples."""
if file_path is None:
return self.default_checkpoint_callback
callback = self.default_file_checkpoint_callback(file_path)
_ = callback # validates extension and path early
def _mcmc_chain_callback(state: dict) -> None:
samples = state.get("samples")
if samples is None:
raise ValueError("Checkpoint missing samples.")
chain_path = (
f"{self.chain_checkpoint_path}/{self.chain_dataset_name}"
)
with AspireFile(file_path, "a") as h5_file:
if chain_path in h5_file:
del h5_file[chain_path]
samples.save(h5_file, path=chain_path, flat=False)
group = h5_file[chain_path]
if (
hasattr(samples, "chain_shape")
and samples.chain_shape is not None
):
group.attrs["chain_shape"] = np.asarray(
samples.chain_shape, dtype=int
)
iteration = state.get("iteration")
if iteration is not None:
group.attrs["iteration"] = int(iteration)
stage = state.get("meta", {}).get("stage")
if stage is not None:
group.attrs["stage"] = str(stage)
sampler_name = state.get("sampler")
if sampler_name is not None:
group.attrs["sampler"] = str(sampler_name)
self.default_checkpoint_callback(state)
return _mcmc_chain_callback
[docs]
def checkpoint_mcmc_chain(
self,
samples: Samples,
iteration: int | None = None,
checkpoint_callback: Callable[[dict], None] | None = None,
checkpoint_every: int | None = None,
checkpoint_file_path: str | None = None,
) -> None:
"""Save an MCMC chain checkpoint."""
if checkpoint_every is not None and checkpoint_every <= 0:
return
if checkpoint_callback is None and checkpoint_file_path is not None:
checkpoint_callback = (
self.default_mcmc_chain_file_checkpoint_callback(
checkpoint_file_path
)
)
if checkpoint_callback is None:
return
state = self.build_checkpoint_state(
samples=samples, iteration=iteration, meta={"stage": "mcmc_chain"}
)
checkpoint_callback(state)
[docs]
class NumpyMCMCSampler(MCMCSampler):
"""MCMCSampler that maps samples and log probabilities to NumPy arrays for
compatibility with numpy-only samplers
"""
[docs]
def log_prob(self, z):
return to_numpy(super().log_prob(z))
[docs]
class Emcee(NumpyMCMCSampler):
@track_calls
[docs]
def sample(
self,
n_samples: int = None,
nwalkers: int = None,
nsteps: int = 500,
rng=None,
discard=0,
checkpoint_callback: Callable[[dict], None] | None = None,
checkpoint_every: int | None = None,
checkpoint_file_path: str | None = None,
**kwargs,
) -> Samples:
from emcee import EnsembleSampler
nwalkers = nwalkers or n_samples
self.sampler = EnsembleSampler(
nwalkers,
self.dims,
log_prob_fn=self.log_prob,
vectorize=True,
)
rng = rng or self.rng or np.random.default_rng()
samples = self.draw_initial_samples(nwalkers)
p0 = samples.x
z0 = to_numpy(self.preconditioning_transform.fit(p0))
self.sampler.run_mcmc(z0, nsteps, **kwargs)
chain = self.sampler.get_chain(discard=discard)
# Transform chain back to original space
chain_z = chain.reshape(-1, self.dims)
chain_x, log_jacobian = self.preconditioning_transform.inverse(chain_z)
chain_x = chain_x.reshape(chain.shape)
# Create MCMCSamples
samples_mcmc = MCMCSamples.from_chain(
chain=chain_x,
parameters=self.parameters,
xp=self.xp,
dtype=self.dtype,
burn_in=discard,
)
self.checkpoint_mcmc_chain(
samples=samples_mcmc,
iteration=nsteps,
checkpoint_callback=checkpoint_callback,
checkpoint_every=checkpoint_every,
checkpoint_file_path=checkpoint_file_path,
)
if n_samples is not None:
logger.info(
f"Subsampling MCMC samples to {n_samples} samples after burn-in."
)
samples_mcmc = samples_mcmc[:n_samples]
return samples_mcmc
[docs]
class MiniPCN(MCMCSampler):
@track_calls
[docs]
def sample(
self,
n_samples: int | None = None,
n_walkers: int | None = None,
rng=None,
target_acceptance_rate=0.234,
n_steps=100,
thin=1,
burnin=0,
last_step_only=False,
step_fn="tpcn",
checkpoint_callback: Callable[[dict], None] | None = None,
checkpoint_every: int | None = None,
checkpoint_file_path: str | None = None,
):
from minipcn import Sampler
from orng import ArrayRNG
rng = rng or self.rng or ArrayRNG(backend=self.backend_str)
n_walkers = n_walkers or n_samples
p0 = self.draw_initial_samples(n_walkers).x
z0 = self.preconditioning_transform.fit(p0)
self.sampler = Sampler(
log_prob_fn=self.log_prob,
step_fn=step_fn,
rng=rng,
dims=self.dims,
target_acceptance_rate=target_acceptance_rate,
xp=self.xp,
)
chain, history = self.sampler.sample(z0, n_steps=n_steps)
_ = history
# Transform the full chain back to original space once so checkpoints
# always capture pre-burn/pre-thin samples.
chain_z_full = chain.reshape(-1, self.dims)
chain_x_full, _ = self.preconditioning_transform.inverse(chain_z_full)
chain_x_full = chain_x_full.reshape(chain.shape)
full_chain_samples = MCMCSamples.from_chain(
chain=chain_x_full,
parameters=self.parameters,
xp=self.xp,
dtype=self.dtype,
thin=1,
burn_in=0,
)
self.checkpoint_mcmc_chain(
samples=full_chain_samples,
iteration=n_steps,
checkpoint_callback=checkpoint_callback,
checkpoint_every=checkpoint_every,
checkpoint_file_path=checkpoint_file_path,
)
if last_step_only:
x = chain_x_full[-1]
samples_mcmc = Samples(
x, xp=self.xp, dtype=self.dtype, parameters=self.parameters
)
samples_mcmc.log_prior = samples_mcmc.array_to_namespace(
self.log_prior(samples_mcmc)
)
samples_mcmc.log_likelihood = samples_mcmc.array_to_namespace(
self.log_likelihood(samples_mcmc)
)
else:
# Apply burn-in and thinning to the full chain for returned samples.
samples_mcmc = full_chain_samples.post_process(
burn_in=burnin, thin=thin
)
if n_samples is not None:
logger.info(
f"Subsampling MCMC samples to {n_samples} samples after burn-in and thinning."
)
samples_mcmc = samples_mcmc[:n_samples]
return samples_mcmc
[docs]
class ParallelTemperedMCMCSampler(MCMCSampler):
"""Wrapper for Parallel Tempered MCMC Samplers"""
[docs]
def log_likelihood_wrapper(self, z):
"""Wrapper for log-likelihood that takes array inputs."""
x, _ = self.preconditioning_transform.inverse(z)
samples = Samples(x, xp=self.xp, dtype=self.dtype)
samples.log_prior = self.log_prior(samples)
samples.log_likelihood = self.log_likelihood(samples)
return to_numpy(samples.log_likelihood)
[docs]
def log_prior_wrapper(self, z):
"""Wrapper for log-prior that takes array inputs."""
x, log_abs_det_jacobian = self.preconditioning_transform.inverse(z)
samples = Samples(x, xp=self.xp, dtype=self.dtype)
# Skip Jacobian to avoid double counting in log_prior and
# log_likelihood
return self.log_prior(samples) + log_abs_det_jacobian