tests: refactor "not slow" tests to be not so slow (#1495)

* wip: speed batched sampling test

* tests: refactor density estimator tests for speed

* speed up mnle api tests

Author: janfb

tests/density_estimator_test.py

@@ -23,6 +23,7 @@
     build_mnpe,
     build_nsf,
     build_resnet_flowmatcher,
+    build_score_estimator,
     build_zuko_bpf,
     build_zuko_gf,
     build_zuko_maf,
@@ -52,9 +53,10 @@
     build_zuko_unaf,
 ]
 
-flowmatching_build_functions = [
+diffusion_builders = [
     build_mlp_flowmatcher,
     build_resnet_flowmatcher,
+    build_score_estimator,
 ]
 
 
@@ -136,7 +138,13 @@ def test_shape_handling_utility_for_density_estimator(
 
 
 @pytest.mark.parametrize(
-    "density_estimator_build_fn", model_builders + flowmatching_build_functions
+    "density_estimator_build_fn",
+    [
+        build_nsf,
+        build_zuko_nsf,
+        build_mlp_flowmatcher,
+        build_score_estimator,
+    ],  # just test nflows, zuko and flowmatching
 )
 @pytest.mark.parametrize("input_sample_dim", (1, 2))
 @pytest.mark.parametrize("input_event_shape", ((1,), (4,)))
@@ -241,10 +249,17 @@ def test_correctness_of_density_estimator_log_prob(
 
 
 @pytest.mark.parametrize(
-    "density_estimator_build_fn", model_builders + flowmatching_build_functions
+    "density_estimator_build_fn",
+    [
+        build_nsf,
+        build_zuko_nsf,
+        build_mlp_flowmatcher,
+    ],  # just test nflows, zuko and flowmatching
 )
-@pytest.mark.parametrize("input_event_shape", ((1,), (4,)))
-@pytest.mark.parametrize("condition_event_shape", ((1,), (7,)))
+@pytest.mark.parametrize(
+    "input_event_shape", ((1,), pytest.param((2,), marks=pytest.mark.slow))
+)
+@pytest.mark.parametrize("condition_event_shape", ((1,), (2,)))
 @pytest.mark.parametrize("sample_shape", ((1000,), (500, 2)))
 def test_correctness_of_batched_vs_seperate_sample_and_log_prob(
     density_estimator_build_fn: Callable,
@@ -267,11 +282,17 @@ def test_correctness_of_batched_vs_seperate_sample_and_log_prob(
     samples = density_estimator.sample(sample_shape, condition=condition)
     samples = samples.reshape(-1, batch_dim, *input_event_shape)  # Flat for comp.
 
+    # Flatten sample_shape to (B*E,) if it is (B, E)
+    if len(sample_shape) > 1:
+        flat_sample_shape = (torch.prod(torch.tensor(sample_shape)).item(),)
+    else:
+        flat_sample_shape = sample_shape
+
     samples_separate1 = density_estimator.sample(
-        (1000,), condition=condition[0][None, ...]
+        flat_sample_shape, condition=condition[0][None, ...]
     )
     samples_separate2 = density_estimator.sample(
-        (1000,), condition=condition[1][None, ...]
+        flat_sample_shape, condition=condition[1][None, ...]
     )
 
     # Check if means are approx. same
@@ -310,12 +331,14 @@ def _build_density_estimator_and_tensors(
     """Helper function for all tests that deal with shapes of density
     estimators."""
 
+    batch_size = 1000
     # Use positive random values for continuous dims (log transform)
-    batch_input = torch.rand((1000, *input_event_shape), dtype=torch.float32) * 10.0
+    batch_input = (
+        torch.rand((batch_size, *input_event_shape), dtype=torch.float32) * 10.0
+    )
     # make last dim discrete for mixed density estimators
-    batch_input[:, -1] = torch.randint(0, 4, (1000,))
-    batch_condition = torch.randn((1000, *condition_event_shape))
-
+    batch_input[:, -1] = torch.randint(0, 4, (batch_size,))
+    batch_condition = torch.randn((batch_size, *condition_event_shape))
     if len(condition_event_shape) > 1:
         embedding_net = CNNEmbedding(condition_event_shape, kernel_size=1)
         z_score_y = "structured"
@@ -335,11 +358,16 @@ def _build_density_estimator_and_tensors(
             z_score_y=z_score_y,
         )
     else:
+        embedding_net_kwarg = (
+            dict(embedding_net_y=embedding_net)
+            if "score" in density_estimator_build_fn.__name__
+            else dict(embedding_net=embedding_net)
+        )
         density_estimator = density_estimator_build_fn(
             torch.randn_like(batch_input),
             torch.randn_like(batch_condition),
-            embedding_net=embedding_net,
             z_score_y=z_score_y,
+            **embedding_net_kwarg,
         )
 
     inputs = batch_input[:batch_dim]

tests/mnle_test.py

@@ -92,12 +92,12 @@ def test_mnle_on_device(
 @pytest.mark.parametrize(
     "sampler", (pytest.param("mcmc", marks=pytest.mark.mcmc), "rejection", "vi")
 )
-@pytest.mark.parametrize("flow_model", ("mdn", "maf", "nsf", "zuko_nsf", "zuko_bpf"))
+@pytest.mark.parametrize("flow_model", ("mdn", "nsf", "zuko_nsf"))
 @pytest.mark.parametrize("z_score_theta", ("independent", "none"))
 def test_mnle_api(flow_model: str, sampler, mcmc_params_fast: dict, z_score_theta: str):
     """Test MNLE API."""
     # Generate mixed data.
-    num_simulations = 100
+    num_simulations = 10
     theta = torch.rand(num_simulations, 2)
     x = torch.cat(
         (
@@ -119,7 +119,7 @@ def test_mnle_api(flow_model: str, sampler, mcmc_params_fast: dict, z_score_thet
         embedding_net=theta_embedding,
     )
     trainer = MNLE(density_estimator=density_estimator)
-    trainer.append_simulations(theta, x).train(max_num_epochs=5)
+    trainer.append_simulations(theta, x).train(max_num_epochs=1)
 
     # Test different samplers.
     posterior = trainer.build_posterior(prior=prior, sample_with=sampler)
@@ -132,7 +132,7 @@ def test_mnle_api(flow_model: str, sampler, mcmc_params_fast: dict, z_score_thet
         posterior.sample(
             (1,),
             init_strategy="proposal",
-            method="slice_np_vectorized",
+            method="hmc_pyro",
             **mcmc_params_fast,
         )
 

tests/posterior_nn_test.py

@@ -151,7 +151,7 @@ def test_batched_sample_log_prob_with_different_x(
 
 
 @pytest.mark.mcmc
-@pytest.mark.parametrize("snlre_method", [NLE_A, NRE_A, NRE_B, NRE_C, NPE_C])
+@pytest.mark.parametrize("snlre_method", [NRE_C])  # it's independent of the method
 @pytest.mark.parametrize("x_o_batch_dim", (0, 1, 2))
 @pytest.mark.parametrize("init_strategy", ["proposal", "resample"])
 @pytest.mark.parametrize(
@@ -179,7 +179,7 @@ def test_batched_mcmc_sample_log_prob_with_different_x(
     inference = snlre_method(prior=prior)
     theta = prior.sample((num_simulations,))
     x = simulator(theta)
-    inference.append_simulations(theta, x).train(max_num_epochs=2)
+    inference.append_simulations(theta, x).train(max_num_epochs=1)
 
     x_o = ones(num_dim) if x_o_batch_dim == 0 else ones(x_o_batch_dim, num_dim)