feat: batched sampling for MCMC

sbi/inference/abc/mcabc.py

@@ -176,7 +176,7 @@
             self.x_o = process_x(x_o, self.x_shape)
         else:
             self.x_shape = x[0, 0].shape
-            self.x_o = process_x(x_o, self.x_shape, allow_iid_x=True)
+            self.x_o = process_x(x_o, self.x_shape)
 
         distances = self.distance(self.x_o, x)
 

sbi/inference/abc/smcabc.py

@@ -389,9 +389,7 @@
             self.x_shape = x[0].shape
         else:
             self.x_shape = x[0, 0].shape
-        self.x_o = process_x(
-            x_o, self.x_shape, allow_iid_x=self.distance.requires_iid_data
-        )
+        self.x_o = process_x(x_o, self.x_shape)
 
         distances = self.distance(self.x_o, x)
         sortidx = torch.argsort(distances)

sbi/inference/posteriors/base_posterior.py

@@ -163,19 +163,15 @@ def set_default_x(self, x: Tensor) -> "NeuralPosterior":
         Returns:
             `NeuralPosterior` that will use a default `x` when not explicitly passed.
         """
-        self._x = process_x(
-            x, x_event_shape=None, allow_iid_x=self.potential_fn.allow_iid_x
-        ).to(self._device)
+        self._x = process_x(x, x_event_shape=None).to(self._device)
         self._map = None
         return self
 
     def _x_else_default_x(self, x: Optional[Array]) -> Tensor:
         if x is not None:
             # New x, reset posterior sampler.
             self._posterior_sampler = None
-            return process_x(
-                x, x_event_shape=None, allow_iid_x=self.potential_fn.allow_iid_x
-            )
+            return process_x(x, x_event_shape=None)
         elif self.default_x is None:
             raise ValueError(
                 "Context `x` needed when a default has not been set."

sbi/inference/posteriors/ensemble_posterior.py

@@ -265,9 +265,7 @@
             `EnsemblePosterior` that will use a default `x` when not explicitly
             passed.
         """
-        self._x = process_x(
-            x, x_event_shape=None, allow_iid_x=self.potential_fn.allow_iid_x
-        ).to(self._device)
+        self._x = process_x(x, x_event_shape=None).to(self._device)
 
         for posterior in self.posteriors:
             posterior.set_default_x(x)
@@ -433,7 +431,7 @@
     def set_x(self, x_o: Optional[Tensor]):
         """Check the shape of the observed data and, if valid, set it."""
         if x_o is not None:
-            x_o = process_x(x_o, allow_iid_x=self.allow_iid_x).to(  # type: ignore
+            x_o = process_x(x_o).to(  # type: ignore
                 self.device
             )
         self._x_o = x_o

sbi/inference/posteriors/mcmc_posterior.py

@@ -1,6 +1,7 @@
 # This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
 # under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
 
+from copy import deepcopy
 from functools import partial
 from math import ceil
 from typing import Any, Callable, Dict, Optional, Union
@@ -20,6 +21,7 @@
 
 from sbi.inference.posteriors.base_posterior import NeuralPosterior
 from sbi.inference.potentials.base_potential import BasePotential
+from sbi.neural_nets.density_estimators.shape_handling import reshape_to_batch_event
 from sbi.samplers.mcmc import (
     IterateParameters,
     PyMCSampler,
@@ -321,6 +323,7 @@
                     thin=thin,  # type: ignore
                     warmup_steps=warmup_steps,  # type: ignore
                     vectorized=(method == "slice_np_vectorized"),
+                    interchangeable_chains=True,
                     num_workers=num_workers,
                     show_progress_bars=show_progress_bars,
                 )
@@ -391,11 +394,74 @@
             Samples from the posteriors of shape (*sample_shape, B, *input_shape)
         """
 
-        # See #1176 for a discussion on the implementation of batched sampling.
-        raise NotImplementedError(
-            "Batched sampling is not implemented for MCMC posterior. \
-            Alternatively you can use `sample` in a loop \
-            [posterior.sample(theta, x_o) for x_o in x]."
+        # Replace arguments that were not passed with their default.
+        method = self.method if method is None else method
+        thin = self.thin if thin is None else thin
+        warmup_steps = self.warmup_steps if warmup_steps is None else warmup_steps
+        num_chains = self.num_chains if num_chains is None else num_chains
+        init_strategy = self.init_strategy if init_strategy is None else init_strategy
+        num_workers = self.num_workers if num_workers is None else num_workers
+        mp_context = self.mp_context if mp_context is None else mp_context
+        init_strategy_parameters = (
+            self.init_strategy_parameters
+            if init_strategy_parameters is None
+            else init_strategy_parameters
+        )
+
+        # custom shape handling to make sure to match the batch size of x and theta
+        # without unnecessary combinations.
+        if len(x.shape) == 1:
+            x = x.unsqueeze(0)
+        batch_size = x.shape[0]
+
+        x = reshape_to_batch_event(x, event_shape=x.shape[1:])
+
+        x_ = x.repeat_interleave(num_chains, dim=0)
+
+        self.potential_fn.set_x(x_, interpret_as_iid=False)
+        self.potential_ = self._prepare_potential(method)  # type: ignore
+
+        # For each observation in the batch, we have num_chains independent chains.
+        num_chains_extended = batch_size * num_chains
+        init_strategy_parameters["num_return_samples"] = num_chains_extended
+        initial_params = self._get_initial_params_batched(
+            x,
+            init_strategy,  # type: ignore
+            num_chains,  # type: ignore
+            num_workers,
+            show_progress_bars,
+            **init_strategy_parameters,
+        )
+        # We need num_samples from each posterior in the batch
+        num_samples = torch.Size(sample_shape).numel() * batch_size
+
+        assert (
+            method == "slice_np_vectorized"
+        ), "Batched sampling only supported for vectorized samplers!"
+
+        with torch.set_grad_enabled(False):
+            transformed_samples = self._slice_np_mcmc(
+                num_samples=num_samples,
+                potential_function=self.potential_,
+                initial_params=initial_params,
+                thin=thin,  # type: ignore
+                warmup_steps=warmup_steps,  # type: ignore
+                vectorized=(method == "slice_np_vectorized"),
+                interchangeable_chains=False,
+                num_workers=num_workers,
+                show_progress_bars=show_progress_bars,
+            )
+
+        samples = self.theta_transform.inv(transformed_samples)
+        sample_shape_len = len(sample_shape)
+        # Samples are of shape (num_samples, num_chains_extended, *input_shape)
+        # concatenate all chains the chains per x together and return.
+        return samples.reshape((batch_size, *sample_shape, -1)).permute(  # type: ignore
+            tuple(range(1, sample_shape_len + 1))
+            + (
+                0,
+                -1,
+            )
         )
 
     def _build_mcmc_init_fn(
@@ -509,7 +575,86 @@
             initial_params = torch.cat(
                 [init_fn() for _ in range(num_chains)]  # type: ignore
             )
+            # initial_params = init_fn()
+        return initial_params
+
+    def _get_initial_params_batched(
+        self,
+        x: torch.Tensor,
+        init_strategy: str,
+        num_chains_per_x: int,
+        num_workers: int,
+        show_progress_bars: bool,
+        **kwargs,
+    ) -> Tensor:
+        """Return initial parameters for MCMC obtained with given init strategy.
+
+        Parallelizes across CPU cores only for SIR.
+
+        Args:
+            x: Batch of observations to create different initial parameters for.
+            init_strategy: Specifies the initialization method. Either of
+                [`proposal`|`sir`|`resample`|`latest_sample`].
+            num_chains_per_x: number of MCMC chains for each x, generates initial params
+                for each x
+            num_workers: number of CPU cores for parallization
+            show_progress_bars: whether to show progress bars for SIR init
+            kwargs: Passed on to `_build_mcmc_init_fn`.
+
+        Returns:
+            Tensor: initial parameters, one for each chain
+        """
+
+        potential_ = deepcopy(self.potential_fn)
+        initial_params = []
+        init_fn = self._build_mcmc_init_fn(
+            self.proposal,
+            potential_fn=potential_,
+            transform=self.theta_transform,
+            init_strategy=init_strategy,  # type: ignore
+            **kwargs,
+        )
+        for xi in x:
+            # Build init function
+            potential_.set_x(xi)
 
+            # Parallelize inits for resampling only.
+            if num_workers > 1 and (
+                init_strategy == "resample" or init_strategy == "sir"
+            ):
+
+                def seeded_init_fn(seed):
+                    torch.manual_seed(seed)
+                    return init_fn()
+
+                seeds = torch.randint(high=2**31, size=(num_chains_per_x,))
+
+                # Generate initial params parallelized over num_workers.
+                with tqdm_joblib(
+                    tqdm(
+                        range(num_chains_per_x),  # type: ignore
+                        disable=not show_progress_bars,
+                        desc=f"""Generating {num_chains_per_x} MCMC inits with
+                                {num_workers} workers.""",
+                        total=num_chains_per_x,
+                    )
+                ):
+                    initial_params = (
+                        initial_params
+                        + [
+                            Parallel(n_jobs=num_workers)(
+                                # pyright: ignore[reportArgumentType]
+                                delayed(seeded_init_fn)(seed)
+                                for seed in seeds
+                            )
+                        ][0]
+                    )  # type: ignore
+            else:
+                initial_params = initial_params + [
+                    init_fn() for _ in range(num_chains_per_x)
+                ]  # type: ignore
+
+        initial_params = torch.cat(initial_params)
         return initial_params
 
     def _slice_np_mcmc(
@@ -520,6 +665,7 @@
         thin: int,
         warmup_steps: int,
         vectorized: bool = False,
+        interchangeable_chains=True,
         num_workers: int = 1,
         init_width: Union[float, ndarray] = 0.01,
         show_progress_bars: bool = True,
@@ -534,6 +680,8 @@
             warmup_steps: Initial number of samples to discard.
             vectorized: Whether to use a vectorized implementation of the
                 `SliceSampler`.
+            interchangeable_chains: Whether chains are interchangeable, i.e., whether
+                we can mix samples between chains.
             num_workers: Number of CPU cores to use.
             init_width: Inital width of brackets.
             show_progress_bars: Whether to show a progressbar during sampling;
@@ -550,9 +698,14 @@
         else:
             SliceSamplerMultiChain = SliceSamplerVectorized
 
+        def multi_obs_potential(params):
+            # Params are of shape (num_chains * num_obs, event).
+            all_potentials = potential_function(params)  # Shape: (num_chains, num_obs)
+            return all_potentials.flatten()
+
         posterior_sampler = SliceSamplerMultiChain(
             init_params=tensor2numpy(initial_params),
-            log_prob_fn=potential_function,
+            log_prob_fn=multi_obs_potential,
             num_chains=num_chains,
             thin=thin,
             verbose=show_progress_bars,
@@ -572,8 +725,11 @@
         # Save sample as potential next init (if init_strategy == 'latest_sample').
         self._mcmc_init_params = samples[:, -1, :].reshape(num_chains, dim_samples)
 
-        # Collect samples from all chains.
-        samples = samples.reshape(-1, dim_samples)[:num_samples]
+        # Update: If chains are interchangeable, return concatenated samples. Otherwise
+        # return samples per chain.
+        if interchangeable_chains:
+            # Collect samples from all chains.
+            samples = samples.reshape(-1, dim_samples)[:num_samples]
 
         return samples.type(torch.float32).to(self._device)
 

sbi/inference/potentials/base_potential.py

@@ -36,15 +36,22 @@
         raise NotImplementedError
 
     @property
-    @abstractmethod
-    def allow_iid_x(self) -> bool:
-        raise NotImplementedError
+    def x_is_iid(self) -> bool:
+        """If x has batch dimension greater than 1, whether to intepret the batch as iid
+        samples or batch of data points."""
+        if self._x_is_iid is not None:
+            return self._x_is_iid
+        else:
+            raise ValueError(
+                "No observed data is available. Use `potential_fn.set_x(x_o)`."
+            )
 
-    def set_x(self, x_o: Optional[Tensor]):
+    def set_x(self, x_o: Optional[Tensor], interpret_as_iid: Optional[bool] = True):
         """Check the shape of the observed data and, if valid, set it."""
         if x_o is not None:
-            x_o = process_x(x_o, allow_iid_x=self.allow_iid_x).to(self.device)
+            x_o = process_x(x_o).to(self.device)
         self._x_o = x_o
+        self._x_is_iid = interpret_as_iid
 
     @property
     def x_o(self) -> Tensor:

sbi/inference/potentials/likelihood_based_potential.py

@@ -54,8 +54,6 @@ def likelihood_estimator_based_potential(
 
 
 class LikelihoodBasedPotential(BasePotential):
-    allow_iid_x = True  # type: ignore
-
     def __init__(
         self,
         likelihood_estimator: ConditionalDensityEstimator,
@@ -85,21 +83,35 @@ def __call__(self, theta: Tensor, track_gradients: bool = True) -> Tensor:
 
         Args:
             theta: The parameter set at which to evaluate the potential function.
+            x_is_iid: Whether to interpret the batch dimension of x_o as iid samples.
             track_gradients: Whether to track the gradients.
 
         Returns:
             The potential $\log(p(x_o|\theta)p(\theta))$.
         """
-
-        # Calculate likelihood over trials and in one batch.
-        log_likelihood_trial_sum = _log_likelihoods_over_trials(
-            x=self.x_o,
-            theta=theta.to(self.device),
-            estimator=self.likelihood_estimator,
-            track_gradients=track_gradients,
-        )
-
-        return log_likelihood_trial_sum + self.prior.log_prob(theta)  # type: ignore
+        if self.x_is_iid:
+            # Calculate likelihood over trials and in one batch.
+            log_likelihood_trial_sum = _log_likelihoods_over_trials(
+                x=self.x_o,
+                theta=theta.to(self.device),
+                estimator=self.likelihood_estimator,
+                track_gradients=track_gradients,
+            )
+            return log_likelihood_trial_sum + self.prior.log_prob(theta)  # type: ignore
+        else:
+            theta_batch_size = theta.shape[0]
+            x_batch_size = self.x_o.shape[0]
+            assert (
+                theta_batch_size == x_batch_size
+            ), f"Batch size mismatch: {theta_batch_size} and {x_batch_size}.\
+                When performing batched sampling for multiple `x`, the batch size of\
+                `theta` must match the batch size of `x`."
+            x = self.x_o.unsqueeze(0)
+            with torch.set_grad_enabled(track_gradients):
+                log_likelihood_batches = self.likelihood_estimator.log_prob(
+                    x, condition=theta
+                )
+            return log_likelihood_batches + self.prior.log_prob(theta)  # type: ignore
 
 
 def _log_likelihoods_over_trials(
@@ -198,7 +210,9 @@ def __init__(
     ):
         super().__init__(likelihood_estimator, prior, x_o, device)
 
-    def __call__(self, theta: Tensor, track_gradients: bool = True) -> Tensor:
+    def __call__(
+        self, theta: Tensor, x_is_iid: bool = True, track_gradients: bool = True
+    ) -> Tensor:
         prior_log_prob = self.prior.log_prob(theta)  # type: ignore
 
         # Shape of `x` is (iid_dim, *event_shape)