aspire-inference 0.1.0a2__py3-none-any.whl

aspire/samplers/smc/blackjax.py

@@ -0,0 +1,330 @@
+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,
+        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,
+            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)
+
+        # 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)
+        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

aspire/samplers/smc/emcee.py

@@ -0,0 +1,75 @@
+import copy
+import logging
+
+import numpy as np
+
+from ...samples import SMCSamples
+from ...utils import track_calls
+from .base import NumpySMCSampler
+
+logger = logging.getLogger(__name__)
+
+
+class EmceeSMC(NumpySMCSampler):
+    @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,
+        sampler_kwargs: dict | None = None,
+        n_final_samples: int | None = None,
+    ):
+        self.sampler_kwargs = sampler_kwargs or {}
+        self.sampler_kwargs.setdefault("nsteps", 5 * self.dims)
+        self.sampler_kwargs.setdefault("progress", True)
+        self.emcee_moves = self.sampler_kwargs.pop("moves", None)
+        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):
+        import emcee
+
+        logger.info("Mutating particles")
+        sampler = emcee.EnsembleSampler(
+            len(particles.x),
+            self.dims,
+            self.log_prob,
+            args=(beta,),
+            vectorize=True,
+            moves=self.emcee_moves,
+        )
+        z = self.fit_preconditioning_transform(particles.x)
+        kwargs = copy.deepcopy(self.sampler_kwargs)
+        if n_steps is not None:
+            kwargs["nsteps"] = n_steps
+        sampler.run_mcmc(z, **kwargs)
+        self.history.mcmc_acceptance.append(
+            np.mean(sampler.acceptance_fraction)
+        )
+        self.history.mcmc_autocorr.append(
+            sampler.get_autocorr_time(
+                quiet=True, discard=int(0.2 * self.sampler_kwargs["nsteps"])
+            )
+        )
+        z = sampler.get_chain(flat=False)[-1, ...]
+        x = self.preconditioning_transform.inverse(z)[0]
+        samples = SMCSamples(x, xp=self.xp, beta=beta)
+        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

aspire/samplers/smc/minipcn.py

@@ -0,0 +1,82 @@
+from functools import partial
+
+import numpy as np
+
+from ...samples import SMCSamples
+from ...utils import to_numpy, track_calls
+from .base import NumpySMCSampler
+
+
+class MiniPCNSMC(NumpySMCSampler):
+    """MiniPCN SMC sampler."""
+
+    rng = None
+
+    def log_prob(self, x, beta=None):
+        return to_numpy(super().log_prob(x, beta))
+
+    @track_calls
+    def sample(
+        self,
+        n_samples: int,
+        n_steps: int = None,
+        min_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,
+    ):
+        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.rng = rng or np.random.default_rng()
+        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_step=min_step,
+            max_n_steps=max_n_steps,
+        )
+
+    def mutate(self, particles, beta, n_steps=None):
+        from minipcn import Sampler
+
+        log_prob_fn = partial(self.log_prob, beta=beta)
+
+        sampler = Sampler(
+            log_prob_fn=log_prob_fn,
+            step_fn=self.sampler_kwargs["step_fn"],
+            rng=self.rng,
+            dims=self.dims,
+            target_acceptance_rate=self.sampler_kwargs[
+                "target_acceptance_rate"
+            ],
+        )
+        # Map to transformed dimension for sampling
+        z = to_numpy(self.fit_preconditioning_transform(particles.x))
+        chain, history = sampler.sample(
+            z,
+            n_steps=n_steps or self.sampler_kwargs["n_steps"],
+        )
+        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)
+        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