add reject_outside_prior_support logic and update rejection sampling tests

Author: satwiksps

sbi/inference/posteriors/direct_posterior.py

@@ -139,6 +139,8 @@ def sample(
         max_sampling_batch_size: int = 10_000,
         sample_with: Optional[str] = None,
         show_progress_bars: bool = True,
+        reject_outside_prior: bool = True,
+        max_sampling_time: Optional[float] = None,
     ) -> Tensor:
         r"""Return samples from posterior distribution $p(\theta|x)$.
 
@@ -149,6 +151,11 @@ def sample(
             sample_with: This argument only exists to keep backward-compatibility with
                 `sbi` v0.17.2 or older. If it is set, we instantly raise an error.
             show_progress_bars: Whether to show sampling progress monitor.
+            reject_outside_prior: If True (default), rejection sampling is used to
+                ensure samples lie within the prior support. If False, samples are drawn
+                directly from the proposal without rejection sampling.
+            max_sampling_time: Optional maximum allowed sampling time in seconds.
+                If exceeded, sampling is aborted and a RuntimeError is raised.
         """
         num_samples = torch.Size(sample_shape).numel()
         x = self._x_else_default_x(x)
@@ -177,15 +184,24 @@ def sample(
                 f"`.build_posterior(sample_with={sample_with}).`"
             )
 
-        samples = rejection.accept_reject_sample(
-            proposal=self.posterior_estimator.sample,
-            accept_reject_fn=lambda theta: within_support(self.prior, theta),
-            num_samples=num_samples,
-            show_progress_bars=show_progress_bars,
-            max_sampling_batch_size=max_sampling_batch_size,
-            proposal_sampling_kwargs={"condition": x},
-            alternative_method="build_posterior(..., sample_with='mcmc')",
-        )[0]  # [0] to return only samples, not acceptance probabilities.
+        if reject_outside_prior:
+            # normal rejection behaviour
+            samples = rejection.accept_reject_sample(
+                proposal=self.posterior_estimator.sample,
+                accept_reject_fn=lambda theta: within_support(self.prior, theta),
+                num_samples=num_samples,
+                show_progress_bars=show_progress_bars,
+                max_sampling_batch_size=max_sampling_batch_size,
+                proposal_sampling_kwargs={"condition": x},
+                alternative_method="build_posterior(..., sample_with='mcmc')",
+                max_sampling_time=max_sampling_time,
+            )[0]
+        else:
+            # bypass rejection sampling entirely
+            samples = self.posterior_estimator.sample(
+                (num_samples,),
+                condition=x,
+            )
 
         return samples[:, 0]  # Remove batch dimension.
 
@@ -195,6 +211,8 @@ def sample_batched(
         x: Tensor,
         max_sampling_batch_size: int = 10_000,
         show_progress_bars: bool = True,
+        reject_outside_prior: bool = True,
+        max_sampling_time: Optional[float] = None,
     ) -> Tensor:
         r"""Given a batch of observations [x_1, ..., x_B] this function samples from
         posteriors $p(\theta|x_1)$, ... ,$p(\theta|x_B)$, in a batched (i.e. vectorized)
@@ -207,6 +225,11 @@ def sample_batched(
                 `batch_dim` corresponds to the number of observations to be drawn.
             max_sampling_batch_size: Maximum batch size for rejection sampling.
             show_progress_bars: Whether to show sampling progress monitor.
+            reject_outside_prior: If True (default), rejection sampling is used to
+                ensure samples lie within the prior support. If False, samples are drawn
+                directly from the proposal without rejection sampling.
+            max_sampling_time: Optional maximum allowed sampling time in seconds.
+                If exceeded, sampling is aborted and a RuntimeError is raised.
 
         Returns:
             Samples from the posteriors of shape (*sample_shape, B, *input_shape)
@@ -242,15 +265,24 @@ def sample_batched(
             )
             max_sampling_batch_size = capped
 
-        samples = rejection.accept_reject_sample(
-            proposal=self.posterior_estimator.sample,
-            accept_reject_fn=lambda theta: within_support(self.prior, theta),
-            num_samples=num_samples,
-            show_progress_bars=show_progress_bars,
-            max_sampling_batch_size=max_sampling_batch_size,
-            proposal_sampling_kwargs={"condition": x},
-            alternative_method="build_posterior(..., sample_with='mcmc')",
-        )[0]
+        if reject_outside_prior:
+            # normal rejection behaviour
+            samples = rejection.accept_reject_sample(
+                proposal=self.posterior_estimator.sample,
+                accept_reject_fn=lambda theta: within_support(self.prior, theta),
+                num_samples=num_samples,
+                show_progress_bars=show_progress_bars,
+                max_sampling_batch_size=max_sampling_batch_size,
+                proposal_sampling_kwargs={"condition": x},
+                alternative_method="build_posterior(..., sample_with='mcmc')",
+                max_sampling_time=max_sampling_time,
+            )[0]
+        else:
+            # bypass rejection sampling entirely
+            samples = self.posterior_estimator.sample(
+                (num_samples,),
+                condition=x,
+            )
 
         return samples
 

sbi/inference/posteriors/vector_field_posterior.py

@@ -158,6 +158,8 @@ def sample(
         max_sampling_batch_size: int = 10_000,
         sample_with: Optional[str] = None,
         show_progress_bars: bool = True,
+        reject_outside_prior: bool = True,
+        max_sampling_time: Optional[float] = None,
     ) -> Tensor:
         r"""Return samples from posterior distribution $p(\theta|x)$.
 
@@ -195,6 +197,11 @@ def sample(
                 use the 'sde' sampling method, the vector field estimator must support
                 it and have the SCORE_DEFINED class attribute set to True.
             show_progress_bars: Whether to show a progress bar during sampling.
+            reject_outside_prior: If True (default), rejection sampling is used to
+                ensure samples lie within the prior support. If False, samples are drawn
+                directly from the proposal without rejection sampling.
+            max_sampling_time: Optional maximum allowed sampling time in seconds.
+                If exceeded, sampling is aborted and a RuntimeError is raised.
         """
 
         if sample_with is None:
@@ -213,13 +220,17 @@ def sample(
         num_samples = torch.Size(sample_shape).numel()
 
         if sample_with == "ode":
-            samples, _ = rejection.accept_reject_sample(
-                proposal=self.sample_via_ode,
-                accept_reject_fn=lambda theta: within_support(self.prior, theta),
-                num_samples=num_samples,
-                show_progress_bars=show_progress_bars,
-                max_sampling_batch_size=max_sampling_batch_size,
-            )
+            if reject_outside_prior:
+                samples, _ = rejection.accept_reject_sample(
+                    proposal=self.sample_via_ode,
+                    accept_reject_fn=lambda theta: within_support(self.prior, theta),
+                    num_samples=num_samples,
+                    show_progress_bars=show_progress_bars,
+                    max_sampling_batch_size=max_sampling_batch_size,
+                    max_sampling_time=max_sampling_time,
+                )
+            else:
+                samples = self.sample_via_ode((num_samples,))
         elif sample_with == "sde":
             proposal_sampling_kwargs = {
                 "predictor": predictor,
@@ -231,14 +242,21 @@ def sample(
                 "max_sampling_batch_size": max_sampling_batch_size,
                 "show_progress_bars": show_progress_bars,
             }
-            samples, _ = rejection.accept_reject_sample(
-                proposal=self._sample_via_diffusion,
-                accept_reject_fn=lambda theta: within_support(self.prior, theta),
-                num_samples=num_samples,
-                show_progress_bars=show_progress_bars,
-                max_sampling_batch_size=max_sampling_batch_size,
-                proposal_sampling_kwargs=proposal_sampling_kwargs,
-            )
+            if reject_outside_prior:
+                samples, _ = rejection.accept_reject_sample(
+                    proposal=self._sample_via_diffusion,
+                    accept_reject_fn=lambda theta: within_support(self.prior, theta),
+                    num_samples=num_samples,
+                    show_progress_bars=show_progress_bars,
+                    max_sampling_batch_size=max_sampling_batch_size,
+                    proposal_sampling_kwargs=proposal_sampling_kwargs,
+                    max_sampling_time=max_sampling_time,
+                )
+            else:
+                samples = self._sample_via_diffusion(
+                    (num_samples,),
+                    **proposal_sampling_kwargs,
+                )
         else:
             raise ValueError(
                 f"Expected sample_with to be 'ode' or 'sde', but got {sample_with}."
@@ -427,6 +445,8 @@ def sample_batched(
         ts: Optional[Tensor] = None,
         max_sampling_batch_size: int = 10000,
         show_progress_bars: bool = True,
+        reject_outside_prior: bool = True,
+        max_sampling_time: Optional[float] = None,
     ) -> Tensor:
         r"""Given a batch of observations [x_1, ..., x_B] this function samples from
         posteriors $p(\theta|x_1)$, ... ,$p(\theta|x_B)$, in a batched (i.e. vectorized)
@@ -449,6 +469,11 @@ def sample_batched(
                 linear grid between t_max and t_min is used.
             max_sampling_batch_size: Maximum batch size for sampling.
             show_progress_bars: Whether to show sampling progress monitor.
+            reject_outside_prior: If True (default), rejection sampling is used to
+                ensure samples lie within the prior support. If False, samples are drawn
+                directly from the proposal without rejection sampling.
+            max_sampling_time: Optional maximum allowed sampling time in seconds.
+                If exceeded, sampling is aborted and a RuntimeError is raised.
 
         Returns:
             Samples from the posteriors of shape (*sample_shape, B, *input_shape)
@@ -477,14 +502,18 @@ def sample_batched(
             max_sampling_batch_size = capped
 
         if self.sample_with == "ode":
-            samples, _ = rejection.accept_reject_sample(
-                proposal=self.sample_via_ode,
-                accept_reject_fn=lambda theta: within_support(self.prior, theta),
-                num_samples=num_samples,
-                num_xos=batch_size,
-                show_progress_bars=show_progress_bars,
-                max_sampling_batch_size=max_sampling_batch_size,
-            )
+            if reject_outside_prior:
+                samples, _ = rejection.accept_reject_sample(
+                    proposal=self.sample_via_ode,
+                    accept_reject_fn=lambda theta: within_support(self.prior, theta),
+                    num_samples=num_samples,
+                    num_xos=batch_size,
+                    show_progress_bars=show_progress_bars,
+                    max_sampling_batch_size=max_sampling_batch_size,
+                    max_sampling_time=max_sampling_time,
+                )
+            else:
+                samples = self.sample_via_ode((num_samples,))
             samples = samples.reshape(
                 sample_shape + batch_shape + self.vector_field_estimator.input_shape
             )
@@ -499,15 +528,21 @@ def sample_batched(
                 "max_sampling_batch_size": max_sampling_batch_size,
                 "show_progress_bars": show_progress_bars,
             }
-            samples, _ = rejection.accept_reject_sample(
-                proposal=self._sample_via_diffusion,
-                accept_reject_fn=lambda theta: within_support(self.prior, theta),
-                num_samples=num_samples,
-                num_xos=batch_size,
-                show_progress_bars=show_progress_bars,
-                max_sampling_batch_size=max_sampling_batch_size,
-                proposal_sampling_kwargs=proposal_sampling_kwargs,
-            )
+            if reject_outside_prior:
+                samples, _ = rejection.accept_reject_sample(
+                    proposal=self._sample_via_diffusion,
+                    accept_reject_fn=lambda theta: within_support(self.prior, theta),
+                    num_samples=num_samples,
+                    num_xos=batch_size,
+                    show_progress_bars=show_progress_bars,
+                    max_sampling_batch_size=max_sampling_batch_size,
+                    proposal_sampling_kwargs=proposal_sampling_kwargs,
+                    max_sampling_time=max_sampling_time,
+                )
+            else:
+                samples = self._sample_via_diffusion(
+                    (num_samples,), **proposal_sampling_kwargs
+                )
             samples = samples.reshape(
                 sample_shape + batch_shape + self.vector_field_estimator.input_shape
             )

sbi/samplers/rejection/rejection.py

@@ -146,8 +146,10 @@ def log_prob(self, theta: Tensor, **kwargs) -> Tensor:
                 raise RuntimeError(
                     "Sampling aborted early because rejection sampling exceeded "
                     "max_sampling_time. This is likely due to extremely low "
-                    "acceptance. Consider switching to MCMC or VI, or checking "
-                    "for model misspecification."
+                    "acceptance. You can disable rejection sampling using "
+                    "`reject_outside_prior=False` to draw samples directly from "
+                    "the trained estimator. Consider switching to MCMC or VI, or "
+                    "checking for model misspecification."
                 )
 
             # Sample and reject.
@@ -319,8 +321,10 @@ def accept_reject_sample(
             raise RuntimeError(
                 "Sampling aborted early because rejection sampling exceeded "
                 "max_sampling_time. This is likely due to extremely low "
-                "acceptance. Consider switching to MCMC or VI, or checking "
-                "for model misspecification."
+                "acceptance. You can disable rejection sampling using "
+                "`reject_outside_prior=False` to draw samples directly from "
+                "the trained estimator. Consider switching to MCMC or VI, or "
+                "checking for model misspecification."
             )
 
         # Sample and reject.