fix: remove NPE rejection prior guard, extend tests for sample_with across all methods (#1839)

* Enhance parameterized tests for various sampling methods

* remove prior validation for rejection sampling in PosteriorEstimatorTrainer

* Remove xfail test for rejection sampling

Author: Dev-Sudarshan

sbi/inference/trainers/npe/npe_base.py

@@ -437,10 +437,6 @@ def build_posterior(
             (the returned log-probability is unnormalized).
         """
 
-        self._check_prior_for_rejection_sampling(
-            prior, sample_with, posterior_parameters
-        )
-
         return super().build_posterior(
             density_estimator,
             prior,
@@ -484,44 +480,6 @@ def _get_potential_function(
         )
         return potential_fn, theta_transform
 
-    def _check_prior_for_rejection_sampling(
-        self,
-        prior: Optional[Distribution],
-        sample_with: Literal["mcmc", "rejection", "vi", "importance", "direct"],
-        posterior_parameters: Optional[
-            Union[
-                DirectPosteriorParameters,
-                MCMCPosteriorParameters,
-                VIPosteriorParameters,
-                RejectionPosteriorParameters,
-                ImportanceSamplingPosteriorParameters,
-            ]
-        ],
-    ) -> None:
-        """
-        Validates that when using rejection sampling, a prior distribution
-        is explicitly provided.
-
-        Args:
-            prior: Prior distribution.
-            sample_with: The sampling method used. Must be one of
-                "mcmc", "rejection", "vi", "importance", or "direct".
-            posterior_parameters: Configuration for building the posterior.
-        """
-
-        if (
-            sample_with == "rejection"
-            or isinstance(posterior_parameters, RejectionPosteriorParameters)
-        ) and prior is None:
-            raise ValueError(
-                "You indicated sampling via rejection sampling but "
-                "haven't passed a prior. As of sbi v0.23.0, you either have"
-                " to pass a prior to perform rejection sampling using the prior"
-                " as proposal, or to use the posterior as proposal, you have to"
-                " use a DirectPosterior via `sample_with='direct' or"
-                " `posterior_parameters=DirectPosteriorParameters`."
-            )
-
     def _loss(
         self,
         theta: Tensor,

tests/posterior_nn_test.py

@@ -11,7 +11,6 @@
 from sbi.inference import (
     FMPE,
     NLE_A,
-    NPE,
     NPE_A,
     NPE_C,
     NPSE,
@@ -23,7 +22,6 @@
 )
 from sbi.inference.posteriors.posterior_parameters import (
     MCMCPosteriorParameters,
-    RejectionPosteriorParameters,
 )
 from sbi.inference.potentials.posterior_based_potential import (
     posterior_estimator_based_potential,
@@ -459,27 +457,6 @@ def simulator(theta):
     inference.build_posterior(density_estimator=nn.Module())
 
 
-@pytest.mark.xfail(
-    raises=ValueError,
-    reason="Prior must be passed through build_posterior method for rejection"
-    " sampling in NPE",
-)
-def test_build_posterior_raises_error_for_rejection_sampling():
-    def simulator(theta):
-        return 1.0 + theta + torch.randn(theta.shape, device=theta.device) * 0.1
-
-    num_dim = 3
-    prior = BoxUniform(low=-2 * torch.ones(num_dim), high=2 * torch.ones(num_dim))
-    theta = prior.sample((300,))
-    x = simulator(theta)
-
-    inference = NPE(prior=prior)
-    inference.append_simulations(theta, x)
-
-    inference.train(max_num_epochs=1)
-    inference.build_posterior(posterior_parameters=RejectionPosteriorParameters())
-
-
 def get_multidim_simulator_embedding(num_dim: int = 5, embedding_dim: int = 10):
     """Returns a simulator producing 2D matrix observations and a CNN embedding net."""
 

tests/posterior_parameters_test.py

@@ -7,7 +7,7 @@
 import pytest
 import torch
 
-from sbi.inference import NPE, NRE
+from sbi.inference import FMPE, NLE, NPE, NPSE, NRE
 from sbi.inference.posteriors.direct_posterior import DirectPosterior
 from sbi.inference.posteriors.importance_posterior import ImportanceSamplingPosterior
 from sbi.inference.posteriors.mcmc_posterior import MCMCPosterior
@@ -416,6 +416,7 @@ def test_if_warning_raised_for_deprecated_build_posterior_parameters(
         get_inference.build_posterior(**params)
 
 
+@pytest.mark.parametrize("inference_class", [NPE, NRE, NLE])
 @pytest.mark.parametrize(
     "sample_with",
     ["mcmc", "vi", "rejection", "importance"],
@@ -424,27 +425,68 @@ def test_if_warning_raised_for_deprecated_build_posterior_parameters(
     raises=ValueError,
     reason="Prior required for non-direct sampling methods",
 )
-def test_resolve_prior_missing_prior_npe(sample_with):
-    inference = NPE()  # no prior
+def test_resolve_prior_missing_prior(inference_class, sample_with):
+    """
+    Test that an error is raised when trying to build a posterior with a sampling method
+    that requires a prior, but no prior is provided to the inference method.
+    """
+    inference = inference_class()  # no prior
 
-    theta = torch.randn(10, 2)
+    theta = torch.randn(50, 2)
     x = theta + 0.1 * torch.randn_like(theta)
 
     inference.append_simulations(theta, x).train(max_num_epochs=1)
 
     inference.build_posterior(sample_with=sample_with)
 
 
-@pytest.mark.parametrize("sample_with", ["direct"])
-def test_resolve_prior_missing_prior_npe_direct_ok(sample_with):
-    inference = NPE()  # no prior
+VALID_SAMPLE_WITH_AND_INFERENCE_METHOD = [
+    (NPE, "direct"),
+    (NPE, "mcmc"),
+    (NPE, "vi"),
+    (NPE, "rejection"),
+    (NPE, "importance"),
+    (NLE, "mcmc"),
+    (NLE, "vi"),
+    (NLE, "rejection"),
+    (NLE, "importance"),
+    (NRE, "mcmc"),
+    (NRE, "vi"),
+    (NRE, "rejection"),
+    (NRE, "importance"),
+    (NPSE, "sde"),
+    (NPSE, "ode"),
+    (FMPE, "sde"),
+    (FMPE, "ode"),
+]
 
-    theta = torch.randn(10, 2)
-    x = theta + 0.1 * torch.randn_like(theta)
+
+@pytest.mark.parametrize(
+    "inference_method, sample_with", VALID_SAMPLE_WITH_AND_INFERENCE_METHOD
+)
+def test_inference_method_with_valid_sample_with(inference_method, sample_with):
+    """
+    Test that the inference method works correctly with valid sample_with options.
+    """
+    num_samples = 2
+    dim_theta = 2
+    num_simulations = 50
+
+    prior = BoxUniform(low=-2 * torch.ones(dim_theta), high=2 * torch.ones(dim_theta))
+    inference = inference_method(prior=prior)
+
+    theta = prior.sample((num_simulations,))
+    x = theta + torch.randn_like(theta) * 0.1
 
     inference.append_simulations(theta, x).train(max_num_epochs=1)
 
     posterior = inference.build_posterior(sample_with=sample_with)
 
-    samples = posterior.sample((5,), x=x[:1])
-    assert samples.shape[0] == 5
+    if sample_with == "vi":
+        posterior.set_default_x(x[0])
+        posterior.train(max_num_iters=5, show_progress_bar=False, quality_control=False)
+        samples = posterior.sample((num_samples,))
+    else:
+        samples = posterior.sample((num_samples,), x=x[0])
+
+    assert samples.shape == torch.Size([num_samples, dim_theta])