from functools import partial
from typing import Any
import numpy as np
from ...samples import SMCSamples
from ...utils import (
asarray,
track_calls,
)
from .base import SMCSampler
[docs]
class MiniPCNSMC(SMCSampler):
"""MiniPCN SMC sampler."""
[docs]
def log_prob(self, x: SMCSamples, beta: float | None = None) -> Any:
return super().log_prob(x, beta)
@track_calls
[docs]
def sample(
self,
n_samples: int,
n_steps: int = None,
min_beta_step: float | None = None,
max_beta_step: float | None = None,
max_n_steps: int | None = None,
adaptive: bool = True,
target_efficiency: float = 0.5,
target_efficiency_rate: float = 1.0,
n_final_samples: int | None = None,
sampler_kwargs: dict | None = None,
rng: np.random.Generator | None = None,
checkpoint_callback=None,
checkpoint_every: int | None = None,
checkpoint_file_path: str | None = None,
resume_from: str | bytes | dict | None = None,
beta_tolerance: float = 1e-6,
store_sample_history: bool = True,
) -> SMCSamples:
from orng import ArrayRNG
self.sampler_kwargs = sampler_kwargs or {}
self.sampler_kwargs.setdefault("n_steps", 5 * self.dims)
self.sampler_kwargs.setdefault("target_acceptance_rate", 0.234)
self.sampler_kwargs.setdefault("step_fn", "tpcn")
self.sampler_kwargs.setdefault("verbose", True)
self.rng = rng or self.rng or ArrayRNG(backend=self.backend_str)
return super().sample(
n_samples,
n_steps=n_steps,
adaptive=adaptive,
target_efficiency=target_efficiency,
target_efficiency_rate=target_efficiency_rate,
n_final_samples=n_final_samples,
min_beta_step=min_beta_step,
max_beta_step=max_beta_step,
max_n_steps=max_n_steps,
checkpoint_callback=checkpoint_callback,
checkpoint_every=checkpoint_every,
checkpoint_file_path=checkpoint_file_path,
resume_from=resume_from,
beta_tolerance=beta_tolerance,
store_sample_history=store_sample_history,
)
[docs]
def mutate(
self, particles: SMCSamples, beta: float, n_steps: int | None = None
) -> SMCSamples:
"""Mutate particles using the MiniPCN sampler.
Parameters
----------
particles : SMCSamples
The current particles to be mutated.
beta : float
The current inverse temperature.
n_steps : int, optional
The number of MCMC steps to take. If None, uses the default from
:code:`sampler_kwargs`.
Returns
-------
SMCSamples
The mutated particles.
"""
from minipcn import Sampler
log_prob_fn = partial(self.log_prob, beta=beta)
kwargs = self.sampler_kwargs.copy()
n_steps_default = kwargs.pop("n_steps")
n_steps = n_steps or n_steps_default
verbose = kwargs.pop("verbose")
sampler = Sampler(
log_prob_fn=log_prob_fn,
rng=self.rng,
dims=self.dims,
xp=self.xp,
**kwargs,
)
# Map to transformed dimension for sampling
z = asarray(
self.fit_preconditioning_transform(particles.x),
xp=self.xp,
dtype=self.dtype,
)
chain, history = sampler.sample(
z,
n_steps=n_steps,
verbose=verbose,
)
x = self.preconditioning_transform.inverse(chain[-1])[0]
self.history.mcmc_acceptance.append(np.mean(history.acceptance_rate))
samples = SMCSamples(
x,
xp=self.xp,
beta=beta,
dtype=self.dtype,
parameters=self.parameters,
)
samples.log_q = samples.array_to_namespace(
self.prior_flow.log_prob(samples.x)
)
samples.log_prior = samples.array_to_namespace(self.log_prior(samples))
samples.log_likelihood = samples.array_to_namespace(
self.log_likelihood(samples)
)
if samples.xp.isnan(samples.log_q).any():
raise ValueError("Log proposal contains NaN values")
return samples