refactor: remove simulate_for_sbi from tests. (#1200)

* refactor: remove simulate_for_sbi from tests.

* fix docstring

sbi/inference/base.py

@@ -593,7 +593,9 @@ def simulate_for_sbi(
         simulator: A function that takes parameters $\theta$ and maps them to
             simulations, or observations, `x`, $\text{sim}(\theta)\to x$. Any
             regular Python callable (i.e. function or class with `__call__` method)
-            can be used.
+            can be used. Note that the simulator should be able to handle numpy
+            arrays for efficient parallelization. You can use
+            `process_simulator` to ensure this.
         proposal: Probability distribution that the parameters $\theta$ are sampled
             from.
         num_simulations: Number of simulations that are run.

tests/inference_on_device_test.py

@@ -249,6 +249,7 @@ def test_validate_theta_and_x_device(training_device: str, data_device: str) ->
     )
 
 
+@pytest.mark.gpu
 @pytest.mark.parametrize(
     "inference_method", [SNPE_A, SNPE_C, SNRE_A, SNRE_B, SNRE_C, SNLE]
 )

tests/linearGaussian_mdn_test.py

@@ -14,13 +14,11 @@
     DirectPosterior,
     MCMCPosterior,
     likelihood_estimator_based_potential,
-    simulate_for_sbi,
 )
 from sbi.simulators.linear_gaussian import (
     linear_gaussian,
     true_posterior_linear_gaussian_mvn_prior,
 )
-from sbi.utils.user_input_checks import process_prior, process_simulator
 from tests.test_utils import check_c2st
 
 
@@ -50,9 +48,9 @@ def simulator(theta: Tensor) -> Tensor:
 
     inference = method(density_estimator="mdn")
 
-    prior, _, prior_returns_numpy = process_prior(prior)
-    simulator = process_simulator(simulator, prior, prior_returns_numpy)
-    theta, x = simulate_for_sbi(simulator, prior, num_simulations)
+    theta = prior.sample((num_simulations,))
+    x = simulator(theta)
+
     estimator = inference.append_simulations(theta, x).train()
     if method == SNPE:
         posterior = DirectPosterior(posterior_estimator=estimator, prior=prior)
@@ -84,6 +82,7 @@ def test_mdn_with_1D_uniform_prior():
     """
     num_dim = 1
     x_o = torch.tensor([[1.0]])
+    num_simulations = 100
     num_samples = 100
 
     # likelihood_mean will be likelihood_shift+theta
@@ -97,9 +96,9 @@ def simulator(theta: Tensor) -> Tensor:
 
     inference = SNPE(density_estimator="mdn")
 
-    prior, _, prior_returns_numpy = process_prior(prior)
-    simulator = process_simulator(simulator, prior, prior_returns_numpy)
-    theta, x = simulate_for_sbi(simulator, prior, 100)
+    theta = prior.sample((num_simulations,))
+    x = simulator(theta)
+
     posterior_estimator = inference.append_simulations(theta, x).train()
     posterior = DirectPosterior(posterior_estimator=posterior_estimator, prior=prior)
     samples = posterior.sample((num_samples,), x=x_o)

tests/linearGaussian_snle_test.py

@@ -15,7 +15,6 @@
     RejectionPosterior,
     VIPosterior,
     likelihood_estimator_based_potential,
-    simulate_for_sbi,
 )
 from sbi.neural_nets import likelihood_nn
 from sbi.simulators.linear_gaussian import (
@@ -26,10 +25,7 @@
     true_posterior_linear_gaussian_mvn_prior,
 )
 from sbi.utils import BoxUniform
-from sbi.utils.user_input_checks import (
-    process_prior,
-    process_simulator,
-)
+from sbi.utils.user_input_checks import process_prior
 
 from .test_utils import check_c2st, get_prob_outside_uniform_prior
 
@@ -42,7 +38,7 @@ def test_api_snle_multiple_trials_and_rounds_map(
     """Test SNLE API with 2 rounds, different priors num trials and MAP."""
     num_rounds = 2
     num_samples = 1
-    num_simulations = 100
+    num_simulations_per_round = 100
 
     if prior_str == "gaussian":
         prior_mean = zeros(num_dim)
@@ -52,18 +48,12 @@ def test_api_snle_multiple_trials_and_rounds_map(
         prior = BoxUniform(-2.0 * ones(num_dim), 2.0 * ones(num_dim))
 
     simulator = diagonal_linear_gaussian
-    prior, _, prior_returns_numpy = process_prior(prior)
-    simulator = process_simulator(simulator, prior, prior_returns_numpy)
     inference = SNLE(prior=prior, density_estimator="mdn", show_progress_bars=False)
 
     proposals = [prior]
     for _ in range(num_rounds):
-        theta, x = simulate_for_sbi(
-            simulator,
-            proposals[-1],
-            num_simulations,
-            simulation_batch_size=num_simulations,
-        )
+        theta = proposals[-1].sample((num_simulations_per_round,))
+        x = simulator(theta)
         inference.append_simulations(theta, x).train(
             training_batch_size=100, max_num_epochs=2
         )
@@ -88,8 +78,8 @@ def test_c2st_snl_on_linear_gaussian_different_dims(
     discard_dims = theta_dim - x_dim
 
     x_o = zeros(1, x_dim)
-    num_samples = 500
-    num_simulations = 1000
+    num_samples = 1000
+    num_simulations = 2000
 
     # likelihood_mean will be likelihood_shift+theta
     likelihood_shift = -1.0 * ones(x_dim)
@@ -119,15 +109,9 @@ def simulator(theta):
     density_estimator = likelihood_nn(model=model_str, num_transforms=3)
     inference = SNLE(density_estimator=density_estimator, show_progress_bars=False)
 
-    prior, _, prior_returns_numpy = process_prior(prior)
-    simulator = process_simulator(simulator, prior, prior_returns_numpy)
-    theta, x = simulate_for_sbi(
-        simulator,
-        prior,
-        num_simulations,
-        simulation_batch_size=num_simulations,
-        seed=1,
-    )
+    theta = prior.sample((num_simulations,))
+    x = simulator(theta)
+
     likelihood_estimator = inference.append_simulations(theta, x).train()
     potential_fn, theta_transform = likelihood_estimator_based_potential(
         prior=prior, likelihood_estimator=likelihood_estimator, x_o=x_o
@@ -181,9 +165,9 @@ def simulator(theta):
     density_estimator = likelihood_nn(model_str, num_transforms=3)
     inference = SNLE(density_estimator=density_estimator, show_progress_bars=False)
 
-    theta, x = simulate_for_sbi(
-        simulator, prior, num_simulations, simulation_batch_size=num_simulations
-    )
+    theta = prior.sample((num_simulations,))
+    x = simulator(theta)
+
     likelihood_estimator = inference.append_simulations(theta, x).train()
 
     # Test inference amortized over trials.
@@ -307,9 +291,9 @@ def simulator(theta):
 
     inference = SNLE(show_progress_bars=False)
 
-    theta, x = simulate_for_sbi(
-        simulator, prior, num_simulations_per_round, simulation_batch_size=50
-    )
+    theta = prior.sample((num_simulations_per_round,))
+    x = simulator(theta)
+
     likelihood_estimator = inference.append_simulations(theta, x).train()
     potential_fn, theta_transform = likelihood_estimator_based_potential(
         prior=prior, likelihood_estimator=likelihood_estimator, x_o=x_o
@@ -322,12 +306,9 @@ def simulator(theta):
         **mcmc_params_accurate,
     )
 
-    theta, x = simulate_for_sbi(
-        simulator,
-        posterior1,
-        num_simulations_per_round,
-        simulation_batch_size=50,
-    )
+    theta = posterior1.sample((num_simulations_per_round,))
+    x = simulator(theta)
+
     likelihood_estimator = inference.append_simulations(theta, x).train()
     potential_fn, theta_transform = likelihood_estimator_based_potential(
         prior=prior, likelihood_estimator=likelihood_estimator, x_o=x_o
@@ -373,9 +354,9 @@ def simulator(theta):
 
     inference = SNLE(show_progress_bars=False)
 
-    theta, x = simulate_for_sbi(
-        simulator, prior, num_simulations_per_round, simulation_batch_size=50
-    )
+    theta = prior.sample((num_simulations_per_round,))
+    x = simulator(theta)
+
     likelihood_estimator = inference.append_simulations(theta, x).train()
     potential_fn, theta_transform = likelihood_estimator_based_potential(
         prior=prior, likelihood_estimator=likelihood_estimator, x_o=x_o
@@ -386,12 +367,9 @@ def simulator(theta):
     )
     posterior1.train()
 
-    theta, x = simulate_for_sbi(
-        simulator,
-        posterior1,
-        num_simulations_per_round,
-        simulation_batch_size=50,
-    )
+    theta = posterior1.sample((num_simulations_per_round,))
+    x = simulator(theta)
+
     likelihood_estimator = inference.append_simulations(theta, x).train()
     potential_fn, theta_transform = likelihood_estimator_based_potential(
         prior=prior, likelihood_estimator=likelihood_estimator, x_o=x_o
@@ -471,9 +449,8 @@ def test_api_snl_sampling_methods(
 
     inference = SNLE(show_progress_bars=False)
 
-    theta, x = simulate_for_sbi(
-        simulator, prior, num_simulations, simulation_batch_size=1000
-    )
+    theta = prior.sample((num_simulations,))
+    x = simulator(theta)
     likelihood_estimator = inference.append_simulations(theta, x).train(
         max_num_epochs=5
     )

tests/linearGaussian_snpe_test.py

@@ -31,7 +31,6 @@
     true_posterior_linear_gaussian_mvn_prior,
 )
 from sbi.utils import RestrictedPrior, get_density_thresholder
-from sbi.utils.user_input_checks import process_prior, process_simulator
 
 from .sbiutils_test import conditional_of_mvn
 from .test_utils import (
@@ -81,12 +80,9 @@ def simulator(theta):
 
     inference = snpe_method(prior, show_progress_bars=False)
 
-    prior, _, prior_returns_numpy = process_prior(prior)
-    simulator = process_simulator(simulator, prior, prior_returns_numpy)
+    theta = prior.sample((num_simulations,))
+    x = simulator(theta)
 
-    theta, x = simulate_for_sbi(
-        simulator, prior, num_simulations, simulation_batch_size=1000
-    )
     posterior_estimator = inference.append_simulations(theta, x).train(
         training_batch_size=100
     )
@@ -238,12 +234,8 @@ def simulator(theta):
         show_progress_bars=False,
     )
 
-    # type: ignore
-    prior, _, prior_returns_numpy = process_prior(prior)
-    simulator = process_simulator(simulator, prior, prior_returns_numpy)
-    theta, x = simulate_for_sbi(
-        simulator, prior, num_simulations, simulation_batch_size=num_simulations
-    )
+    theta = prior.sample((num_simulations,))
+    x = simulator(theta)
 
     inference = inference.append_simulations(theta, x)
     posterior_estimator = inference.train(
@@ -473,6 +465,7 @@ def test_api_force_first_round_loss(
 
     num_dim = 2
     x_o = zeros(1, num_dim)
+    num_simulations = 1000
 
     # likelihood_mean will be likelihood_shift+theta
     likelihood_shift = -1.0 * ones(num_dim)
@@ -483,13 +476,12 @@ def simulator(theta):
         return linear_gaussian(theta, likelihood_shift, likelihood_cov)
 
     inference = SNPE_C(prior, show_progress_bars=False)
-    prior, _, prior_returns_numpy = process_prior(prior)
-    simulator = process_simulator(simulator, prior, prior_returns_numpy)
 
     proposal = prior
     for _ in range(2):
         train_proposal = proposal if pass_proposal_to_append else None
-        theta, x = simulate_for_sbi(simulator, proposal, 1000)
+        theta = proposal.sample((num_simulations,))
+        x = simulator(theta)
         _ = inference.append_simulations(theta, x, proposal=train_proposal).train(
             force_first_round_loss=force_first_round_loss, max_num_epochs=2
         )
@@ -662,12 +654,9 @@ def simulator(theta):
 
     inference = SNPE_C(density_estimator="mdn", show_progress_bars=False)
 
-    prior, _, prior_returns_numpy = process_prior(prior)
-    simulator = process_simulator(simulator, prior, prior_returns_numpy)
+    theta = prior.sample((num_simulations,))
+    x = simulator(theta)
 
-    theta, x = simulate_for_sbi(
-        simulator, prior, num_simulations, simulation_batch_size=1000
-    )
     posterior_mdn = inference.append_simulations(theta, x).train(
         training_batch_size=100
     )
@@ -703,16 +692,8 @@ def simulator(theta):
         return linear_gaussian(theta, likelihood_shift, likelihood_cov)
 
     inference = snpe_method(prior, show_progress_bars=False)
-    prior, _, prior_returns_numpy = process_prior(prior)
-    simulator = process_simulator(simulator, prior, prior_returns_numpy)
-
-    theta, x = simulate_for_sbi(
-        simulator,
-        prior,
-        num_simulations,
-        simulation_batch_size=num_simulations,
-        num_workers=6,
-    )
+    theta = prior.sample((num_simulations,))
+    x = simulator(theta)
     posterior_estimator = inference.append_simulations(theta, x).train(
         max_num_epochs=2, **extra_kwargs
     )

tests/linearGaussian_snre_test.py

@@ -29,7 +29,6 @@
     samples_true_posterior_linear_gaussian_uniform_prior,
     true_posterior_linear_gaussian_mvn_prior,
 )
-from sbi.utils.user_input_checks import process_prior, process_simulator
 from tests.test_utils import (
     check_c2st,
     get_dkl_gaussian_prior,
@@ -54,17 +53,10 @@ def test_api_snre_multiple_trials_and_rounds_map(
     simulator = diagonal_linear_gaussian
     inference = snre_method(prior=prior, classifier="mlp", show_progress_bars=False)
 
-    prior, _, prior_returns_numpy = process_prior(prior)
-    simulator = process_simulator(simulator, prior, prior_returns_numpy)
-
     proposals = [prior]
     for _ in range(num_rounds):
-        theta, x = simulate_for_sbi(
-            simulator,
-            proposals[-1],
-            num_simulations,
-            simulation_batch_size=num_simulations,
-        )
+        theta = proposals[-1].sample((num_simulations,))
+        x = simulator(theta)
         inference.append_simulations(theta, x).train(
             training_batch_size=100, max_num_epochs=2
         )
@@ -113,12 +105,8 @@ def simulator(theta):
 
     inference = snre_method(classifier="resnet", show_progress_bars=False)
 
-    prior, _, prior_returns_numpy = process_prior(prior)
-    simulator = process_simulator(simulator, prior, prior_returns_numpy)
-
-    theta, x = simulate_for_sbi(
-        simulator, prior, num_simulations, simulation_batch_size=100
-    )
+    theta = prior.sample((num_simulations,))
+    x = simulator(theta)
     ratio_estimator = inference.append_simulations(theta, x).train()
 
     x_o = zeros(1, x_dim)