Score-based iid sampling

sbi/inference/posteriors/score_posterior.py

@@ -105,6 +105,8 @@ def sample(
         corrector_params: Optional[Dict] = None,
         steps: int = 500,
         ts: Optional[Tensor] = None,
+        iid_method: str = "auto_gauss",
+        iid_params: Optional[Dict] = None,
         max_sampling_batch_size: int = 10_000,
         sample_with: Optional[str] = None,
         show_progress_bars: bool = True,
@@ -123,6 +125,9 @@ def sample(
             steps: Number of steps to take for the Euler-Maruyama method.
             ts: Time points at which to evaluate the diffusion process. If None, a
                 linear grid between t_max and t_min is used.
+            iid_method: Which method to use for computing the score in the iid setting.
+                We currently support "fnpe", "gauss", "auto_gauss", "jac_gauss".
+            iid_params: Additional parameters passed to the iid method.
             max_sampling_batch_size: Maximum batch size for sampling.
             sample_with: Deprecated - use `.build_posterior(sample_with=...)` prior to
                 `.sample()`.
@@ -138,7 +143,10 @@ def sample(
 
         x = self._x_else_default_x(x)
         x = reshape_to_batch_event(x, self.score_estimator.condition_shape)
-        self.potential_fn.set_x(x, x_is_iid=True)
+        is_iid = x.ndim > 1 and x.shape[0] > 1
+        self.potential_fn.set_x(
+            x, x_is_iid=is_iid, iid_method=iid_method, iid_params=iid_params
+        )
 
         num_samples = torch.Size(sample_shape).numel()
 
@@ -176,6 +184,7 @@ def sample(
     def _sample_via_diffusion(
         self,
         sample_shape: Shape = torch.Size(),
+        x: Optional[Tensor] = None,
         predictor: Union[str, Predictor] = "euler_maruyama",
         corrector: Optional[Union[str, Corrector]] = None,
         predictor_params: Optional[Dict] = None,
@@ -244,6 +253,7 @@ def _sample_via_diffusion(
     def sample_via_ode(
         self,
         sample_shape: Shape = torch.Size(),
+        x: Optional[Tensor] = None,
     ) -> Tensor:
         r"""Return samples from posterior distribution with probability flow ODE.
 

sbi/inference/potentials/score_based_potential.py

@@ -11,6 +11,7 @@
 from zuko.transforms import FreeFormJacobianTransform
 
 from sbi.inference.potentials.base_potential import BasePotential
+from sbi.inference.potentials.score_fn_iid import get_iid_method
 from sbi.neural_nets.estimators.score_estimator import ConditionalScoreEstimator
 from sbi.neural_nets.estimators.shape_handling import (
     reshape_to_batch_event,
@@ -57,7 +58,8 @@ def __init__(
         score_estimator: ConditionalScoreEstimator,
         prior: Optional[Distribution],
         x_o: Optional[Tensor] = None,
-        iid_method: str = "iid_bridge",
+        iid_method: str = "auto_gauss",
+        iid_params: Optional[dict] = None,
         device: str = "cpu",
     ):
         r"""Returns the score function for score-based methods.
@@ -66,30 +68,37 @@ def __init__(
             score_estimator: The neural network modelling the score.
             prior: The prior distribution.
             x_o: The observed data at which to evaluate the posterior.
-            iid_method: Which method to use for computing the score. Currently, only
-                `iid_bridge` as proposed in Geffner et al. is implemented.
+            iid_method: Which method to use for computing the score in the iid setting.
+                We currently support "fnpe", "gauss", "auto_gauss", "jac_gauss".
+            iid_params: Parameters for the iid method, for arguments see ScoreFnIID.
             device: The device on which to evaluate the potential.
         """
         self.score_estimator = score_estimator
         self.score_estimator.eval()
         self.iid_method = iid_method
+        self.iid_params = iid_params
         super().__init__(prior, x_o, device=device)
 
     def set_x(
         self,
         x_o: Optional[Tensor],
         x_is_iid: Optional[bool] = False,
+        iid_method: str = "auto_gauss",
+        iid_params: Optional[dict] = None,
         rebuild_flow: Optional[bool] = True,
     ):
         """
         Set the observed data and whether it is IID.
         Args:
-        x_o: The observed data.
-        x_is_iid: Whether the observed data is IID (if batch_dim>1).
-        rebuild_flow: Whether to save (overwrrite) a low-tolerance flow model, useful if
-        the flow needs to be evaluated many times (e.g. for MAP calculation).
+            x_o: The observed data.
+            x_is_iid: Whether the observed data is IID (if batch_dim>1).
+            rebuild_flow: Whether to save (overwrrite) a low-tolerance flow model,
+                useful if the flow needs to be evaluated many times
+                (e.g. for MAP calculation).
         """
         super().set_x(x_o, x_is_iid)
+        self.iid_method = iid_method
+        self.iid_params = iid_params
         if rebuild_flow and self._x_o is not None:
             # By default, we want a high-tolerance flow.
             # This flow will be used mainly for MAP calculations, hence we want to save
@@ -172,10 +181,16 @@ def gradient(
                     input=theta, condition=self.x_o, time=time
                 )
             else:
-                raise NotImplementedError(
-                    "Score accumulation for IID data is not yet implemented."
+                assert self.prior is not None, "Prior is required for iid methods."
+
+                method_iid = get_iid_method(self.iid_method)
+                # Always creating a new object every call is not efficient...
+                score_fn_iid = method_iid(
+                    self.score_estimator, self.prior, **(self.iid_params or {})
                 )
 
+                score = score_fn_iid(theta, self.x_o, time)  # type: ignore
+
         return score
 
     def get_continuous_normalizing_flow(
@@ -217,9 +232,6 @@ def rebuild_flow(
         x_density_estimator = reshape_to_batch_event(
             self.x_o, event_shape=self.score_estimator.condition_shape
         )
-        assert x_density_estimator.shape[0] == 1, (
-            "PosteriorScoreBasedPotential supports only x batchsize of 1`."
-        )
 
         flow = self.get_continuous_normalizing_flow(
             condition=x_density_estimator, atol=atol, rtol=rtol, exact=exact