updated tests for new api

tests/inference_with_NaN_simulator_test.py

@@ -1,20 +1,17 @@
 # This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
 # under the Affero General Public License v3, see <https://www.gnu.org/licenses/>.
 
-from sbi.utils.sbiutils import handle_invalid_x
 import pytest
 import torch
-
-from sbi.inference import SNPE_C, SRE, SNL, prepare_for_sbi
-from torch import zeros, ones, eye
+from torch import eye, ones, zeros
 
 import sbi.utils as utils
-
+from sbi.inference import SNL, SNPE_C, SRE, prepare_for_sbi
 from sbi.simulators.linear_gaussian import (
-    samples_true_posterior_linear_gaussian_uniform_prior,
     linear_gaussian,
+    samples_true_posterior_linear_gaussian_uniform_prior,
 )
-
+from sbi.utils.sbiutils import handle_invalid_x
 from tests.test_utils import check_c2st
 
 
@@ -76,9 +73,7 @@ def linear_gaussian_nan(
     infer = method(*prepare_for_sbi(linear_gaussian_nan, prior))
 
     posterior = infer(
-        num_rounds=1,
-        num_simulations_per_round=num_simulations,
-        exclude_invalid_x=exclude_invalid_x,
+        num_simulations=num_simulations, exclude_invalid_x=exclude_invalid_x,
     ).set_default_x(x_o)
 
     samples = posterior.sample((num_samples,))

tests/linearGaussian_mdn.py

@@ -48,9 +48,7 @@ def simulator(theta: Tensor) -> Tensor:
 
     infer = method(simulator, prior, density_estimator="mdn")
 
-    posterior = infer(
-        num_rounds=1, num_simulations_per_round=1000, training_batch_size=50
-    ).set_default_x(x_o)
+    posterior = infer(num_simulations=1000, training_batch_size=50).set_default_x(x_o)
     samples = posterior.sample((num_samples,))
 
     # Compute the c2st and assert it is near chance level of 0.5.

tests/linearGaussian_snle_test.py

@@ -16,8 +16,8 @@
     samples_true_posterior_linear_gaussian_uniform_prior,
     true_posterior_linear_gaussian_mvn_prior,
 )
-from tests.test_utils import check_c2st, get_prob_outside_uniform_prior
 from sbi.utils.torchutils import configure_default_device
+from tests.test_utils import check_c2st, get_prob_outside_uniform_prior
 
 
 @pytest.mark.parametrize("num_dim", (1, 3))
@@ -44,9 +44,9 @@ def test_api_snl_on_linearGaussian(num_dim: int, set_seed):
         show_progress_bars=False,
     )
 
-    posterior = infer(num_rounds=1, num_simulations_per_round=1000, max_num_epochs=5)
+    posterior = infer(num_simulations=1000, max_num_epochs=5)
 
-    posterior.sample(sample_shape=(num_samples,), x=x_o, mcmc_parameters = {"thin": 3})
+    posterior.sample(sample_shape=(num_samples,), x=x_o, mcmc_parameters={"thin": 3})
 
 
 def test_c2st_snl_on_linearGaussian_different_dims(set_seed):
@@ -97,8 +97,8 @@ def test_c2st_snl_on_linearGaussian_different_dims(set_seed):
         device=device,
     )
 
-    posterior = infer(num_rounds=1, num_simulations_per_round=5000)  # type: ignore
-    samples = posterior.sample((num_samples,), x=x_o, mcmc_parameters = {"thin": 3})
+    posterior = infer(num_simulations=5000)  # type: ignore
+    samples = posterior.sample((num_samples,), x=x_o, mcmc_parameters={"thin": 3})
 
     # Compute the c2st and assert it is near chance level of 0.5.
     check_c2st(samples, target_samples, alg="snle_a")
@@ -146,9 +146,7 @@ def simulator(theta):
 
     infer.provide_presimulated(external_theta, external_x)
 
-    posterior = infer(
-        num_rounds=1, num_simulations_per_round=1000, training_batch_size=100,
-    ).set_default_x(x_o)
+    posterior = infer(num_simulations=1000, training_batch_size=100).set_default_x(x_o)
     samples = posterior.sample((num_samples,))
 
     # Compute the c2st and assert it is near chance level of 0.5.
@@ -196,9 +194,9 @@ def test_c2st_snl_on_linearGaussian(num_dim: int, prior_str: str, set_seed):
         show_progress_bars=False,
     )
 
-    posterior = infer(num_rounds=1, num_simulations_per_round=1000).set_default_x(x_o)
+    posterior = infer(num_simulations=1000).set_default_x(x_o)
 
-    samples = posterior.sample(sample_shape=(num_samples,), mcmc_parameters = {"thin": 3})
+    samples = posterior.sample(sample_shape=(num_samples,), mcmc_parameters={"thin": 3})
 
     # Check performance based on c2st accuracy.
     check_c2st(samples, target_samples, alg=f"snle_a-{prior_str}-prior")
@@ -246,9 +244,10 @@ def test_c2st_multi_round_snl_on_linearGaussian(set_seed):
         show_progress_bars=False,
     )
 
-    posterior = infer(num_rounds=2, x_o=x_o, num_simulations_per_round=500)
+    posterior1 = infer(num_simulations=500).focus_training_on(x_o)
+    posterior = infer(num_simulations=500, proposal=posterior1)
 
-    samples = posterior.sample(sample_shape=(num_samples,), mcmc_parameters = {"thin": 3})
+    samples = posterior.sample(sample_shape=(num_samples,), mcmc_parameters={"thin": 3})
 
     # Check performance based on c2st accuracy.
     check_c2st(samples, target_samples, alg="multi-round-snl")
@@ -283,6 +282,6 @@ def test_api_snl_sampling_methods(mcmc_method: str, prior_str: str, set_seed):
         show_progress_bars=False,
     )
 
-    posterior = infer(num_rounds=1, num_simulations_per_round=200, max_num_epochs=5)
+    posterior = infer(num_simulations=200, max_num_epochs=5)
 
-    posterior.sample(sample_shape=(num_samples,), x=x_o, mcmc_parameters = {"thin": 3})
+    posterior.sample(sample_shape=(num_samples,), x=x_o, mcmc_parameters={"thin": 3})

tests/linearGaussian_snpe_test.py

@@ -2,7 +2,6 @@
 # under the Affero General Public License v3, see <https://www.gnu.org/licenses/>.
 
 from __future__ import annotations
-from sbi.utils.torchutils import configure_default_device
 
 import pytest
 from torch import eye, ones, zeros
@@ -16,6 +15,7 @@
     samples_true_posterior_linear_gaussian_uniform_prior,
     true_posterior_linear_gaussian_mvn_prior,
 )
+from sbi.utils.torchutils import configure_default_device
 from tests.test_utils import (
     check_c2st,
     get_dkl_gaussian_prior,
@@ -70,9 +70,7 @@ def simulator(theta):
         device=device,
     )
 
-    posterior = infer(
-        num_rounds=1, num_simulations_per_round=2000, training_batch_size=100
-    ).set_default_x(x_o)
+    posterior = infer(num_simulations=2000, training_batch_size=100).set_default_x(x_o)
     samples = posterior.sample((num_samples,))
 
     # Compute the c2st and assert it is near chance level of 0.5.
@@ -164,7 +162,7 @@ def simulator(theta):
         sample_with_mcmc=False,
     )
 
-    posterior = infer(num_rounds=1, num_simulations_per_round=2000)  # type: ignore
+    posterior = infer(num_simulations=2000)  # type: ignore
     samples = posterior.sample((num_samples,), x=x_o)
 
     # Compute the c2st and assert it is near chance level of 0.5.
@@ -214,9 +212,7 @@ def simulator(theta):
 
     infer.provide_presimulated(external_theta, external_x)
 
-    posterior = infer(
-        num_rounds=1, num_simulations_per_round=1000, training_batch_size=100,
-    ).set_default_x(x_o)
+    posterior = infer(num_simulations=1000, training_batch_size=100).set_default_x(x_o)
     samples = posterior.sample((num_samples,))
 
     # Compute the c2st and assert it is near chance level of 0.5.
@@ -270,16 +266,19 @@ def simulator(theta):
         density_estimator="maf",
         show_progress_bars=False,
     )
-    call_args = dict(num_rounds=2, x_o=x_o, num_simulations_per_round=1000)
 
     if method_str == "snpe_b":
         infer = SNPE_B(simulation_batch_size=10, **creation_args)
-        posterior = infer(**call_args)
+        posterior1 = infer(num_simulations=1000)
+        posterior1.focus_training_on(x_o)
+        posterior = infer(num_simulations=1000, proposal=posterior1)
     elif method_str == "snpe_c":
         infer = SNPE_C(
             simulation_batch_size=50, sample_with_mcmc=False, **creation_args
         )
-        posterior = infer(num_atoms=10, **call_args)
+        posterior1 = infer(num_simulations=500, num_atoms=10)
+        posterior1.focus_training_on(x_o)
+        posterior = infer(num_simulations=500, num_atoms=10, proposal=posterior1)
 
     samples = posterior.sample((num_samples,))
 
@@ -335,7 +334,7 @@ def simulator(theta):
         show_progress_bars=False,
     )
 
-    posterior = infer(num_rounds=1, num_simulations_per_round=1000, max_num_epochs=5)
+    posterior = infer(num_simulations=1000, max_num_epochs=5)
 
     # Posterior should be corrected for leakage even if num_rounds just 1.
     samples = posterior.sample((10,), x=x_o)
@@ -368,4 +367,4 @@ def simulator(theta):
         sample_with_mcmc=False,
     )
 
-    return infer(num_rounds=1, num_simulations_per_round=1000).set_default_x(x_o)
+    return infer(num_simulations=1000).set_default_x(x_o)

tests/linearGaussian_snre_test.py

@@ -16,12 +16,12 @@
     samples_true_posterior_linear_gaussian_uniform_prior,
     true_posterior_linear_gaussian_mvn_prior,
 )
+from sbi.utils.torchutils import configure_default_device
 from tests.test_utils import (
     check_c2st,
     get_dkl_gaussian_prior,
     get_prob_outside_uniform_prior,
 )
-from sbi.utils.torchutils import configure_default_device
 
 
 @pytest.mark.parametrize("num_dim", (1, 3))
@@ -45,7 +45,7 @@ def test_api_sre_on_linearGaussian(num_dim: int):
         show_progress_bars=False,
     )
 
-    posterior = infer(num_rounds=1, num_simulations_per_round=1000, max_num_epochs=5)
+    posterior = infer(num_simulations=1000, max_num_epochs=5)
 
     posterior.sample(sample_shape=(10,), x=x_o, mcmc_parameters={"num_chains": 2})
 
@@ -100,8 +100,8 @@ def simulator(theta):
         device=device,
     )
 
-    posterior = infer(num_rounds=1, num_simulations_per_round=5000)
-    samples = posterior.sample((num_samples,), x=x_o, mcmc_parameters = {"thin": 3})
+    posterior = infer(num_simulations=5000)
+    samples = posterior.sample((num_samples,), x=x_o, mcmc_parameters={"thin": 3})
 
     # Compute the c2st and assert it is near chance level of 0.5.
     check_c2st(samples, target_samples, alg="snpe_c")
@@ -148,9 +148,7 @@ def simulator(theta):
 
     infer.provide_presimulated(external_theta, external_x)
 
-    posterior = infer(
-        num_rounds=1, num_simulations_per_round=1000, training_batch_size=100,
-    ).set_default_x(x_o)
+    posterior = infer(num_simulations=1000, training_batch_size=100).set_default_x(x_o)
     samples = posterior.sample((num_samples,))
 
     # Compute the c2st and assert it is near chance level of 0.5.
@@ -215,9 +213,9 @@ def simulator(theta):
     infer = SRE(**kwargs) if method_str == "sre" else AALR(**kwargs)
 
     # Should use default `num_atoms=10` for SRE; `num_atoms=2` for AALR
-    posterior = infer(num_rounds=1, num_simulations_per_round=1000).set_default_x(x_o)
+    posterior = infer(num_simulations=1000).set_default_x(x_o)
 
-    samples = posterior.sample(sample_shape=(num_samples,), mcmc_parameters = {"thin": 3})
+    samples = posterior.sample(sample_shape=(num_samples,), mcmc_parameters={"thin": 3})
 
     # Check performance based on c2st accuracy.
     check_c2st(samples, target_samples, alg=f"sre-{prior_str}-{method_str}")
@@ -280,6 +278,6 @@ def test_api_sre_sampling_methods(mcmc_method: str, prior_str: str, set_seed):
         show_progress_bars=False,
     )
 
-    posterior = infer(num_rounds=1, num_simulations_per_round=200, max_num_epochs=5)
+    posterior = infer(num_simulations=200, max_num_epochs=5)
 
     posterior.sample(sample_shape=(10,), x=x_o)

tests/posterior_nn_test.py

@@ -4,7 +4,7 @@
 from __future__ import annotations
 
 import pytest
-from torch import zeros, ones, eye
+from torch import eye, ones, zeros
 from torch.distributions import MultivariateNormal
 
 from sbi.inference import SNPE_C, prepare_for_sbi
@@ -17,7 +17,7 @@ def test_log_prob_with_different_x():
 
     prior = MultivariateNormal(loc=zeros(num_dim), covariance_matrix=eye(num_dim))
     posterior = SNPE_C(*prepare_for_sbi(diagonal_linear_gaussian, prior))(
-        num_rounds=1, num_simulations_per_round=1000
+        num_simulations=1000
     )
 
     _ = posterior.sample((10,), x=ones(1, num_dim))

tests/user_input_checks_test.py

@@ -2,14 +2,14 @@
 # under the Affero General Public License v3, see <https://www.gnu.org/licenses/>.
 
 from __future__ import annotations
-from sbi.utils.get_nn_models import posterior_nn
 
 from typing import Callable
 
 import pytest
 import torch
+from pyknos.mdn.mdn import MultivariateGaussianMDN
 from scipy.stats import beta, multivariate_normal, uniform
-from torch import Tensor, eye, ones, zeros
+from torch import Tensor, eye, nn, ones, zeros
 from torch.distributions import Beta, Distribution, Gamma, MultivariateNormal, Uniform
 
 from sbi.inference import SNPE_C
@@ -26,6 +26,7 @@
     PytorchReturnTypeWrapper,
     ScipyPytorchWrapper,
 )
+from sbi.utils.get_nn_models import posterior_nn
 from sbi.utils.torchutils import BoxUniform
 
 
@@ -273,7 +274,7 @@ def test_inference_with_user_sbi_problems(user_simulator: Callable, user_prior):
     )
 
     # Run inference.
-    _ = infer(num_rounds=1, num_simulations_per_round=100, max_num_epochs=2)
+    _ = infer(num_simulations=100, max_num_epochs=2)
 
 
 @pytest.mark.parametrize(
@@ -375,8 +376,6 @@ def test_invalid_inputs():
         joint.log_prob(ones(10, 4, 1))
 
 
-from torch import nn
-from pyknos.mdn.mdn import MultivariateGaussianMDN
 
 
 @pytest.mark.parametrize(
@@ -429,4 +428,3 @@ def test_passing_custom_density_estimator(arg):
         density_estimator = arg
     prior = MultivariateNormal(torch.zeros(2), torch.eye(2))
     SNPE_C(diagonal_linear_gaussian, prior, density_estimator=density_estimator)
-