aspire-inference 0.1.0a2__py3-none-any.whl

aspire/samplers/mcmc.py

@@ -0,0 +1,158 @@
+import numpy as np
+
+from ..samples import Samples, to_numpy
+from ..utils import track_calls
+from .base import Sampler
+
+
+class MCMCSampler(Sampler):
+    def draw_initial_samples(self, n_samples: int) -> Samples:
+        """Draw initial samples from the prior flow."""
+        # 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:
+            n_to_draw = n_samples - n_samples_drawn
+            x, log_q = self.prior_flow.sample_and_log_prob(n_to_draw)
+            new_samples = Samples(x, xp=self.xp, log_q=log_q)
+            new_samples.log_prior = new_samples.array_to_namespace(
+                self.log_prior(new_samples)
+            )
+            valid = self.xp.isfinite(new_samples.log_prior)
+            n_valid = 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]
+
+        samples.log_likelihood = samples.array_to_namespace(
+            self.log_likelihood(samples)
+        )
+        return samples
+
+    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)
+        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 to_numpy(log_prob).flatten()
+
+
+class Emcee(MCMCSampler):
+    @track_calls
+    def sample(
+        self,
+        n_samples: int,
+        nwalkers: int = None,
+        nsteps: int = 500,
+        rng=None,
+        discard=0,
+        **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 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)
+
+        z = self.sampler.get_chain(flat=True, discard=discard)
+        x = self.preconditioning_transform.inverse(z)[0]
+
+        x_evidence, log_q = self.prior_flow.sample_and_log_prob(n_samples)
+        samples_evidence = Samples(x_evidence, log_q=log_q, xp=self.xp)
+        samples_evidence.log_prior = self.log_prior(samples_evidence)
+        samples_evidence.log_likelihood = self.log_likelihood(samples_evidence)
+        samples_evidence.compute_weights()
+
+        samples_mcmc = Samples(x, xp=self.xp, 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)
+        )
+        samples_mcmc.log_evidence = samples_mcmc.array_to_namespace(
+            samples_evidence.log_evidence
+        )
+        samples_mcmc.log_evidence_error = samples_mcmc.array_to_namespace(
+            samples_evidence.log_evidence_error
+        )
+
+        return samples_mcmc
+
+
+class MiniPCN(MCMCSampler):
+    @track_calls
+    def sample(
+        self,
+        n_samples,
+        rng=None,
+        target_acceptance_rate=0.234,
+        n_steps=100,
+        thin=1,
+        burnin=0,
+        last_step_only=False,
+        step_fn="tpcn",
+    ):
+        from minipcn import Sampler
+
+        rng = rng or np.random.default_rng()
+        p0 = self.draw_initial_samples(n_samples).x
+
+        z0 = to_numpy(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,
+        )
+
+        chain, history = self.sampler.sample(z0, n_steps=n_steps)
+
+        if last_step_only:
+            z = chain[-1]
+        else:
+            z = chain[burnin::thin].reshape(-1, self.dims)
+
+        x = self.preconditioning_transform.inverse(z)[0]
+
+        samples_mcmc = Samples(x, xp=self.xp, 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)
+        )
+        return samples_mcmc

aspire/samplers/smc/base.py

@@ -0,0 +1,312 @@
+import logging
+from typing import Callable
+
+import array_api_compat.numpy as np
+
+from ...flows.base import Flow
+from ...history import SMCHistory
+from ...samples import SMCSamples
+from ...utils import (
+    asarray,
+    effective_sample_size,
+    track_calls,
+    update_at_indices,
+)
+from ..mcmc import MCMCSampler
+
+logger = logging.getLogger(__name__)
+
+
+class SMCSampler(MCMCSampler):
+    """Base class for Sequential Monte Carlo samplers."""
+
+    def __init__(
+        self,
+        log_likelihood: Callable,
+        log_prior: Callable,
+        dims: int,
+        prior_flow: Flow,
+        xp: Callable,
+        parameters: list[str] | None = None,
+        rng: np.random.Generator | 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,
+            parameters=parameters,
+            preconditioning_transform=preconditioning_transform,
+        )
+        self.rng = rng or np.random.default_rng()
+        self._adapative_target_efficiency = False
+
+    @property
+    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._adapative_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._adapative_target_efficiency = True
+
+    def current_target_efficiency(self, beta: float) -> float:
+        """Get the current target efficiency based on beta."""
+        if self._adapative_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
+
+    def determine_beta(
+        self,
+        samples: SMCSamples,
+        beta: float,
+        beta_step: float,
+        min_step: float,
+    ) -> 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_step : float
+            The minimum beta step size.
+
+        Returns
+        -------
+        beta : float
+            The new beta value.
+        min_step : float
+            The new minimum step size if adaptive_min_step is True.
+        """
+        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
+            tol = 1e-5
+            eff_beta_max = effective_sample_size(
+                samples.log_weights(beta_max)
+            ) / len(samples)
+            if eff_beta_max >= self.current_target_efficiency(beta_prev):
+                beta_min = 1.0
+            target_eff = self.current_target_efficiency(beta_prev)
+            while beta_max - beta_min > tol:
+                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 self.adaptive_min_step:
+                min_step = min_step * (1 - beta_prev) / (1 - beta_star)
+            beta = max(beta_star, beta_prev + min_step)
+            beta = min(beta, 1.0)
+        return beta, min_step
+
+    @track_calls
+    def sample(
+        self,
+        n_samples: int,
+        n_steps: int | None = None,
+        adaptive: bool = True,
+        min_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,
+    ) -> SMCSamples:
+        samples = self.draw_initial_samples(n_samples)
+        samples = SMCSamples.from_samples(samples, xp=self.xp, beta=0.0)
+        self.fit_preconditioning_transform(samples.x)
+
+        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.history = SMCHistory()
+
+        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
+        beta = 0.0
+
+        if min_step is None:
+            if max_n_steps is None:
+                min_step = 0.0
+                self.adaptive_min_step = False
+            else:
+                min_step = 1 / max_n_steps
+                self.adaptive_min_step = True
+        else:
+            self.adaptive_min_step = False
+
+        iterations = 0
+        while True:
+            iterations += 1
+
+            beta, min_step = self.determine_beta(
+                samples,
+                beta,
+                beta_step,
+                min_step,
+            )
+            self.history.eff_target.append(
+                self.current_target_efficiency(beta)
+            )
+
+            logger.info(f"it {iterations} - beta: {beta}")
+            self.history.beta.append(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(ess)
+            logger.info(
+                f"it {iterations} - ESS: {ess:.1f} ({eff:.2f} efficiency)"
+            )
+            self.history.ess_target.append(
+                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(log_evidence_ratio)
+            self.history.log_norm_ratio_var.append(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 beta == 1.0 or (
+                max_n_steps is not None and iterations >= max_n_steps
+            ):
+                break
+
+        # If n_final_samples is not None, perform additional mutations steps
+        if n_final_samples is not None:
+            logger.info(f"Generating {n_final_samples} final samples")
+            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))
+        )
+
+        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
+
+    def mutate(self, particles):
+        raise NotImplementedError
+
+    def log_prob(self, z, beta=None):
+        x, log_abs_det_jacobian = self.preconditioning_transform.inverse(z)
+        samples = SMCSamples(x, xp=self.xp)
+        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 = update_at_indices(
+            log_prob, self.xp.isnan(log_prob), -self.xp.inf
+        )
+        return log_prob
+
+
+class NumpySMCSampler(SMCSampler):
+    def __init__(
+        self,
+        log_likelihood,
+        log_prior,
+        dims,
+        prior_flow,
+        xp,
+        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,
+            xp,
+            parameters=parameters,
+            preconditioning_transform=preconditioning_transform,
+        )