aspire-inference 0.1.0a7__py3-none-any.whl

aspire/samplers/smc/base.py

@@ -0,0 +1,318 @@
+import logging
+from typing import Any, 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,
+        dtype: Any | str | None = None,
+        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,
+            dtype=dtype,
+            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, dtype=self.dtype
+        )
+        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, 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 = 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,
+        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,
+        )

aspire/samplers/smc/blackjax.py

@@ -0,0 +1,332 @@
+import logging
+from functools import partial
+
+import numpy as np
+
+from ...samples import SMCSamples
+from ...utils import asarray, to_numpy, track_calls
+from .base import SMCSampler
+
+logger = logging.getLogger(__name__)
+
+
+class BlackJAXSMC(SMCSampler):
+    """BlackJAX SMC sampler."""
+
+    def __init__(
+        self,
+        log_likelihood,
+        log_prior,
+        dims,
+        prior_flow,
+        xp,
+        dtype=None,
+        parameters=None,
+        preconditioning_transform=None,
+        rng: np.random.Generator | None = None,  # New parameter
+    ):
+        # For JAX compatibility, we'll keep the original xp
+        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,
+        )
+        self.key = None
+        self.rng = rng or np.random.default_rng()
+
+    def log_prob(self, x, beta=None):
+        """Log probability function compatible with BlackJAX."""
+        # Convert to original xp format for computation
+        if hasattr(x, "__array__"):
+            x_original = asarray(x, self.xp)
+        else:
+            x_original = x
+
+        # Transform back to parameter space
+        x_params, log_abs_det_jacobian = (
+            self.preconditioning_transform.inverse(x_original)
+        )
+        samples = SMCSamples(x_params, xp=self.xp, dtype=self.dtype)
+
+        # Compute log probabilities
+        log_q = self.prior_flow.log_prob(samples.x)
+        samples.log_q = samples.array_to_namespace(log_q)
+        samples.log_prior = samples.array_to_namespace(self.log_prior(samples))
+        samples.log_likelihood = samples.array_to_namespace(
+            self.log_likelihood(samples)
+        )
+
+        # Compute target log probability
+        log_prob = samples.log_p_t(
+            beta=beta
+        ).flatten() + samples.array_to_namespace(log_abs_det_jacobian)
+
+        # Handle NaN values
+        log_prob = self.xp.where(
+            self.xp.isnan(log_prob), -self.xp.inf, log_prob
+        )
+
+        return log_prob
+
+    @track_calls
+    def sample(
+        self,
+        n_samples: int,
+        n_steps: int = 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_key=None,
+    ):
+        """Sample using BlackJAX SMC.
+
+        Parameters
+        ----------
+        n_samples : int
+            Number of samples to draw.
+        n_steps : int
+            Number of SMC steps.
+        adaptive : bool
+            Whether to use adaptive tempering.
+        target_efficiency : float
+            Target efficiency for adaptive tempering.
+        n_final_samples : int | None
+            Number of final samples to return.
+        sampler_kwargs : dict | None
+            Additional arguments for the BlackJAX sampler.
+            - algorithm: str, one of "nuts", "hmc", "rwmh", "random_walk"
+            - n_steps: int, number of MCMC steps per mutation
+            - step_size: float, step size for HMC/NUTS
+            - inverse_mass_matrix: array, inverse mass matrix
+            - sigma: float or array, proposal covariance for random walk MH
+            - num_integration_steps: int, integration steps for HMC
+        rng_key : jax.random.key| None
+            JAX random key for reproducibility.
+        """
+        self.sampler_kwargs = sampler_kwargs or {}
+        self.sampler_kwargs.setdefault("n_steps", 5 * self.dims)
+        self.sampler_kwargs.setdefault("algorithm", "nuts")
+        self.sampler_kwargs.setdefault("step_size", 1e-3)
+        self.sampler_kwargs.setdefault("inverse_mass_matrix", None)
+        self.sampler_kwargs.setdefault("sigma", 0.1)  # For random walk MH
+
+        # Initialize JAX random key
+        if rng_key is None:
+            import jax
+
+            self.key = jax.random.key(42)
+        else:
+            self.key = rng_key
+
+        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,
+        )
+
+    def mutate(self, particles, beta, n_steps=None):
+        """Mutate particles using BlackJAX MCMC."""
+        import blackjax
+        import jax
+
+        logger.debug("Mutating particles with BlackJAX")
+
+        # Split the random key
+        self.key, subkey = jax.random.split(self.key)
+
+        # Transform particles to latent space
+        z = self.fit_preconditioning_transform(particles.x)
+
+        # Convert to JAX arrays
+        z_jax = jax.numpy.asarray(to_numpy(z))
+
+        # Create log probability function for this beta
+        log_prob_fn = partial(self._jax_log_prob, beta=beta)
+
+        # Choose BlackJAX algorithm
+        algorithm = self.sampler_kwargs["algorithm"].lower()
+
+        n_steps = n_steps or self.sampler_kwargs["n_steps"]
+
+        if algorithm == "rwmh" or algorithm == "random_walk":
+            # Initialize Random Walk Metropolis-Hastings sampler
+            sigma = self.sampler_kwargs.get("sigma", 0.1)
+
+            # BlackJAX RMH expects a transition function, not a covariance
+            if isinstance(sigma, (int, float)):
+                # Create a multivariate normal proposal function
+                def proposal_fn(key, position):
+                    return position + sigma * jax.random.normal(
+                        key, position.shape
+                    )
+            else:
+                # For more complex covariance structures
+                if len(sigma) == self.dims:
+                    # Diagonal covariance
+                    sigma_diag = jax.numpy.array(sigma)
+
+                    def proposal_fn(key, position):
+                        return position + sigma_diag * jax.random.normal(
+                            key, position.shape
+                        )
+                else:
+                    # Full covariance matrix
+                    sigma_matrix = jax.numpy.array(sigma)
+
+                    def proposal_fn(key, position):
+                        return position + jax.random.multivariate_normal(
+                            key, jax.numpy.zeros(self.dims), sigma_matrix
+                        )
+
+            rwmh = blackjax.rmh(log_prob_fn, proposal_fn)
+
+            # Initialize states for each particle
+            n_particles = z_jax.shape[0]
+            keys = jax.random.split(subkey, n_particles)
+
+            # Vectorized initialization and sampling
+            def init_and_sample(key, z_init):
+                state = rwmh.init(z_init)
+
+                def one_step(state, key):
+                    state, info = rwmh.step(key, state)
+                    return state, (state, info)
+
+                keys = jax.random.split(key, n_steps)
+                final_state, (states, infos) = jax.lax.scan(
+                    one_step, state, keys
+                )
+                return final_state, infos
+
+            # Vectorize over particles
+            final_states, all_infos = jax.vmap(init_and_sample)(keys, z_jax)
+
+            # Extract final positions
+            z_final = final_states.position
+
+            # Calculate acceptance rates
+            acceptance_rates = jax.numpy.mean(all_infos.is_accepted, axis=1)
+            mean_acceptance = jax.numpy.mean(acceptance_rates)
+
+        elif algorithm == "nuts":
+            # Initialize step size and mass matrix if not provided
+            inverse_mass_matrix = self.sampler_kwargs["inverse_mass_matrix"]
+            if inverse_mass_matrix is None:
+                inverse_mass_matrix = jax.numpy.eye(self.dims)
+
+            step_size = self.sampler_kwargs["step_size"]
+
+            # Initialize NUTS sampler
+            nuts = blackjax.nuts(
+                log_prob_fn,
+                step_size=step_size,
+                inverse_mass_matrix=inverse_mass_matrix,
+            )
+
+            # Initialize states for each particle
+            n_particles = z_jax.shape[0]
+            keys = jax.random.split(subkey, n_particles)
+
+            # Vectorized initialization and sampling
+            def init_and_sample(key, z_init):
+                state = nuts.init(z_init)
+
+                def one_step(state, key):
+                    state, info = nuts.step(key, state)
+                    return state, (state, info)
+
+                keys = jax.random.split(key, self.sampler_kwargs["n_steps"])
+                final_state, (states, infos) = jax.lax.scan(
+                    one_step, state, keys
+                )
+                return final_state, infos
+
+            # Vectorize over particles
+            final_states, all_infos = jax.vmap(init_and_sample)(keys, z_jax)
+
+            # Extract final positions
+            z_final = final_states.position
+
+            # Calculate acceptance rates
+            acceptance_rates = jax.numpy.mean(all_infos.is_accepted, axis=1)
+            mean_acceptance = jax.numpy.mean(acceptance_rates)
+
+        elif algorithm == "hmc":
+            # Initialize HMC sampler
+            hmc = blackjax.hmc(
+                log_prob_fn,
+                step_size=self.sampler_kwargs["step_size"],
+                num_integration_steps=self.sampler_kwargs.get(
+                    "num_integration_steps", 10
+                ),
+                inverse_mass_matrix=(
+                    self.sampler_kwargs["inverse_mass_matrix"]
+                    or jax.numpy.eye(self.dims)
+                ),
+            )
+
+            # Similar vectorized sampling as NUTS
+            n_particles = z_jax.shape[0]
+            keys = jax.random.split(subkey, n_particles)
+
+            def init_and_sample(key, z_init):
+                state = hmc.init(z_init)
+
+                def one_step(state, key):
+                    state, info = hmc.step(key, state)
+                    return state, (state, info)
+
+                keys = jax.random.split(key, self.sampler_kwargs["n_steps"])
+                final_state, (states, infos) = jax.lax.scan(
+                    one_step, state, keys
+                )
+                return final_state, infos
+
+            final_states, all_infos = jax.vmap(init_and_sample)(keys, z_jax)
+            z_final = final_states.position
+            acceptance_rates = jax.numpy.mean(all_infos.is_accepted, axis=1)
+            mean_acceptance = jax.numpy.mean(acceptance_rates)
+
+        else:
+            raise ValueError(f"Unsupported algorithm: {algorithm}")
+
+        # Convert back to parameter space
+        z_final_np = to_numpy(z_final)
+        x_final = self.preconditioning_transform.inverse(z_final_np)[0]
+
+        # Store MCMC diagnostics
+        self.history.mcmc_acceptance.append(float(mean_acceptance))
+
+        # Create new samples
+        samples = SMCSamples(x_final, xp=self.xp, beta=beta, dtype=self.dtype)
+        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
+
+    def _jax_log_prob(self, z, beta):
+        """JAX-compatible log probability function."""
+        import jax.numpy as jnp
+
+        # Single particle version for JAX
+        z_expanded = jnp.expand_dims(z, 0)  # Add batch dimension
+        log_prob = self.log_prob(z_expanded, beta=beta)
+        return log_prob[0]  # Remove batch dimension