Add timeout handling and reject_outside_prior option to rejection sampling

Author: satwiksps

sbi/inference/posteriors/rejection_posterior.py

@@ -137,6 +137,8 @@ def sample(
         m: Optional[float] = None,
         sample_with: Optional[str] = None,
         show_progress_bars: bool = True,
+        reject_outside_prior: bool = True,
+        max_sampling_time: Optional[float] = None,
     ):
         r"""Return samples from posterior $p(\theta|x)$ via rejection sampling.
 
@@ -147,7 +149,13 @@ 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.
         Returns:
             Samples from posterior.
         """
@@ -180,18 +188,22 @@ def sample(
         )
         m = self.m if m is None else m
 
-        samples, _ = rejection_sample(
-            potential,
-            proposal=self.proposal,
-            num_samples=num_samples,
-            show_progress_bars=show_progress_bars,
-            warn_acceptance=0.01,
-            max_sampling_batch_size=max_sampling_batch_size,
-            num_samples_to_find_max=num_samples_to_find_max,
-            num_iter_to_find_max=num_iter_to_find_max,
-            m=m,
-            device=self._device,
-        )
+        if reject_outside_prior:
+            samples, _ = rejection_sample(
+                potential,
+                proposal=self.proposal,
+                num_samples=num_samples,
+                show_progress_bars=show_progress_bars,
+                warn_acceptance=0.01,
+                max_sampling_batch_size=max_sampling_batch_size,
+                num_samples_to_find_max=num_samples_to_find_max,
+                num_iter_to_find_max=num_iter_to_find_max,
+                m=m,
+                max_sampling_time=max_sampling_time,
+                device=self._device,
+            )
+        else:
+            samples = self.proposal.sample((num_samples,))
 
         return samples.reshape((*sample_shape, -1))
 

sbi/samplers/rejection/rejection.py

@@ -188,8 +188,12 @@ def log_prob(self, theta: Tensor, **kwargs) -> Tensor:
                 logging.warning(
                     f"""Only {acceptance_rate:.3%} proposal samples were accepted. It
                         may take a long time to collect the remaining {num_remaining}
-                        samples. Consider interrupting (Ctrl-C) and switching to a
-                        different sampling method with
+                        samples. You can prevent long runtimes by
+                        setting `max_sampling_time` to limit runtime, or disabling
+                        rejection sampling (e.g. via `reject_outside_prior=False` in
+                        `posterior.sample()` when available).
+                        Alternatively, consider interrupting (Ctrl-C) and switching
+                        to a different sampling method with
                         `build_posterior(..., sample_with='mcmc')`. or
                         `build_posterior(..., sample_with='vi')`."""
                 )
@@ -387,10 +391,13 @@ def accept_reject_sample(
             else:
                 warn_msg = f"""Only {min_acceptance_rate:.3%} proposal samples are
                     accepted. It may take a long time to collect the remaining
-                    {num_remaining} samples. """
+                    {num_remaining} samples. You can prevent very long runtimes by
+                    setting `max_sampling_time` to limit runtime, or disabling
+                    rejection sampling (e.g. via `reject_outside_prior=False` in
+                    `posterior.sample()` when available)."""
                 if alternative_method is not None:
-                    warn_msg += f"""Consider interrupting (Ctrl-C) and switching to
-                    `{alternative_method}`."""
+                    warn_msg += f"""Alternatively, consider interrupting (Ctrl-C)
+                    and switching to `{alternative_method}`."""
                 logging.warning(warn_msg)
 
             leakage_warning_raised = True  # Ensure warning is raised just once.

tests/rejection_timeout_test.py

@@ -4,6 +4,7 @@
 import pytest
 import torch
 
+from sbi.inference.posteriors.rejection_posterior import RejectionPosterior
 from sbi.samplers.rejection import accept_reject_sample, rejection_sample
 
 
@@ -28,12 +29,12 @@ def always_reject_fn(x):
 def test_accept_reject_sample_timeout():
     proposal = DummyProposal()
 
-    with pytest.raises(RuntimeError):
+    with pytest.raises(RuntimeError, match="rejection sampling exceeded"):
         accept_reject_sample(
             proposal=proposal,
             accept_reject_fn=always_reject_fn,
             num_samples=5,
-            max_sampling_time=0.2,
+            max_sampling_time=0.01,
         )
 
 
@@ -44,11 +45,46 @@ def test_rejection_sample_timeout():
     def dummy_potential_fn(x):
         return torch.full((x.shape[0],), -1e6)
 
-    with pytest.raises(RuntimeError):
+    with pytest.raises(RuntimeError, match="rejection sampling exceeded"):
         rejection_sample(
             potential_fn=dummy_potential_fn,
             proposal=proposal,
             num_samples=5,
-            max_sampling_time=0.2,
+            max_sampling_time=0.01,
             m=1e12,
         )
+
+
+@pytest.mark.slow
+def test_rejection_posterior_timeout():
+    prior = torch.distributions.MultivariateNormal(torch.zeros(1), torch.eye(1))
+
+    class DummyPotential:
+        """
+        Minimal compliant potential implementing CustomPotential interface.
+        Forces rejection by returning very low log-probability.
+        """
+
+        device = "cpu"
+
+        def __call__(
+            self, theta: torch.Tensor, x_o: torch.Tensor = None
+        ) -> torch.Tensor:
+            return torch.full((theta.shape[0],), -1e6)
+
+        def set_x(self, x):
+            pass
+
+        def to(self, device):
+            self.device = device
+            return self
+
+    posterior = RejectionPosterior(
+        potential_fn=DummyPotential(),
+        proposal=prior,
+    )
+
+    posterior.set_default_x(torch.zeros(1))
+
+    with pytest.raises(RuntimeError, match="max_sampling_time"):
+        posterior.sample((5,), max_sampling_time=0.01)