Repository: sbi-dev/sbi

Issue:
Warning that simulator does not accept batch
In the HH example, I get the following warning:

![Screen Shot 2020-05-06 at 9 15 36 am](https://user-images.githubusercontent.com/24639769/81146976-a0371280-8f79-11ea-93c8-f8881e23bfb6.png)

So, even though I explicitly set `simulation_batch_size=1`, it still warns. Is this wanted? Or is it just a bug?


Discussion:
the warning is raised during input checks, irrespective of the batch size used later in the code. 
I think this is want we want because we want batchable simulators in sbi.

Do you think we should turn that off when `simulation_batch_size=1` is set?
I see. Yes, I would rather turn off the warning if `simulation_batch_size=1` is set. @meteore @jan-matthis what do you think?
We could add to the warning sth like "if you want to simulate in batches."
Hm. I think that the majority of users with mechanistic models will not write a simulator that takes batches. Instead, they would go for the multiprocessing option (once we provide it). So my choice would be no warning in the case `simulation_batch_size=1`.

Also it's ugly to have a warning in the tutorial notebook ;) And artificially looping over parameter sets in the simulator is not nice either.
But the problem is that the current sbi code might rely on simulators returning shape `(1, theta_numel)` for a single theta. Thus, in the user input checks 
- we still check for batchability of the `user_simulator`, 
- don't wrap it to be a batch simulator with a for loop, 
- don't warn, 
- but wrap it to return a leading batch dimension? 
Yes, to me this sounds good.

Resolving patch:
diff --git a/sbi/user_input/user_input_checks.py b/sbi/user_input/user_input_checks.py
--- a/sbi/user_input/user_input_checks.py
+++ b/sbi/user_input/user_input_checks.py
@@ -373,22 +373,22 @@ def wrap_as_pytorch_simulator(
         f"{type(theta)}"
 
         # Define a wrapper function to PyTorch
-        def pytorch_simulator(theta: Tensor):
+        def pytorch_simulator(theta: Tensor) -> Tensor:
             return torch.as_tensor(simulator(theta.numpy()), dtype=float32)
 
     else:
         # Define a wrapper to make sure that the output of the simulator is float32
-        def pytorch_simulator(theta: Tensor):
+        def pytorch_simulator(theta: Tensor) -> Tensor:
             return torch.as_tensor(simulator(theta), dtype=float32)
 
     return pytorch_simulator
 
 
 def wrap_as_batch_simulator(simulator: Callable, prior) -> Callable:
-    """Return a batch simulator.
+    """Return simulator with batch dimension.
 
-    A batch simulator can handle a batch of parameters and return the
-    corresponding batch of simulated data.
+    Return the unchanged simulator if it can simulate multiple parameter vectors per
+    call. Otherwise, wrap to simulator returning a leading batch dimension of shape [1].
     """
 
     is_batched_simulator = True
@@ -399,32 +399,24 @@ def wrap_as_batch_simulator(simulator: Callable, prior) -> Callable:
         assert simulator_batch_size == batch_size
     except:
         is_batched_simulator = False
-        warnings.warn(
-            """Simulator can't handle batches of parameters. It will loop over batches
-            in Python with `map`. Consider vectorising the simulator natively for
-            performance."""
-        )
-        # make sure the simulator does not return batch dim.
-        single_simulation = simulator(prior.sample())
-        assert (
-            not single_simulation.shape[0] == 1
-        ), f"""The simulator can't handle batches and returns a singleton batch
-            dimension: {single_simulation.shape}. A simulator that can't handle batches
-            of parameters must not return a batch dim."""
 
-    return simulator if is_batched_simulator else get_batched_simulator(simulator)
+    return simulator if is_batched_simulator else get_batch_dim_simulator(simulator)
 
 
-def get_batched_simulator(simulator: Callable) -> Callable:
+def get_batch_dim_simulator(simulator: Callable) -> Callable:
     """Return simulator wrapped with `map` to handle batches of parameters."""
 
-    def batched_simulator(theta: Tensor) -> Tensor:
-        theta = ensure_theta_batched(theta)
-        # XXX: this should be handled with more care, e.g., enable multiprocessing
-        # XXX: with Pool() as p: p.map(...)
-        return torch.stack(list(map(simulator, theta)))
+    def batch_dim_simulator(theta: Tensor) -> Tensor:
+        batch_shape, *_ = theta.shape
+        assert (
+            batch_shape == 1
+        ), f"This simulator can handle one single thetas, theta.shape={theta.shape}"
+        # Remove possible singleton dimensions in the user-simulator.
+        simulation = simulator(theta.squeeze()).squeeze()
+        # Return with leading batch dimennsion.
+        return simulation.unsqueeze(0)
 
-    return batched_simulator
+    return batch_dim_simulator
 
 
 def process_x_o(
@@ -434,7 +426,7 @@ def process_x_o(
 
     Args:
         x_o: observed data as provided by the user.
-        simulator: simulator function as provided by the user.
+        simulator: simulator function as processed with sbi checks.
         prior: prior object with .sample() and .log_prob() methods.
 
     Returns:
@@ -450,9 +442,9 @@ def process_x_o(
 
     # Get unbatched simulated data by sampling from prior and simulator.
     # cast to tensor for comparison
-    simulated_data = torch.as_tensor(simulator(prior.sample()), dtype=float32).squeeze(
-        0
-    )
+    simulated_data = torch.as_tensor(
+        simulator(prior.sample((1,))), dtype=float32
+    ).squeeze(0)
 
     # Get data shape by ommitting the batch dimension.
     x_o_shape = x_o.shape[1:]
@@ -531,7 +523,7 @@ def check_sbi_problem(simulator: Callable, prior, x_o: Tensor):
         prior: prior (Distribution like)
         x_o: observed data
     """
-    num_prior_samples = 2
+    num_prior_samples = 1
     theta = prior.sample((num_prior_samples,))
     theta_batch_shape, *_ = theta.shape
     simulation = simulator(theta)


Test patch:
diff --git a/tests/user_input_checks_test.py b/tests/user_input_checks_test.py
--- a/tests/user_input_checks_test.py
+++ b/tests/user_input_checks_test.py
@@ -197,9 +197,7 @@ def test_process_matrix_observation():
         (linear_gaussian, BoxUniform(zeros(2), ones(2))),
         (numpy_linear_gaussian, UserNumpyUniform(zeros(2), ones(2), True)),
         (linear_gaussian_no_batch, BoxUniform(zeros(2), ones(2))),
-        pytest.param(
-            list_simulator, BoxUniform(zeros(2), ones(2)), marks=pytest.mark.xfail,
-        ),
+        pytest.param(list_simulator, BoxUniform(zeros(2), ones(2)),),
     ),
 )
 def test_process_simulator(simulator: Callable, prior: Distribution):
@@ -207,7 +205,7 @@ def test_process_simulator(simulator: Callable, prior: Distribution):
     prior, theta_dim, prior_returns_numpy = process_prior(prior)
     simulator = process_simulator(simulator, prior, prior_returns_numpy)
 
-    n_batch = 2
+    n_batch = 1
     x = simulator(prior.sample((n_batch,)))
 
     assert isinstance(x, Tensor), "Processed simulator must return Tensor."
@@ -231,24 +229,12 @@ def test_process_simulator(simulator: Callable, prior: Distribution):
             BoxUniform(zeros(3, dtype=torch.float64), ones(3, dtype=torch.float64)),
             zeros(3),
         ),
-        pytest.param(
-            list_simulator,
-            BoxUniform(zeros(3), ones(3)),
-            zeros(1, 3),
-            marks=pytest.mark.xfail,
-        ),
         (
             numpy_linear_gaussian,
             UserNumpyUniform(zeros(3), ones(3), return_numpy=True),
             np.zeros((1, 3)),
         ),
-        pytest.param(  # test simulator always returning batch dim.
-            lambda _: zeros(1, 480),
-            BoxUniform(zeros(2), ones(2)),
-            zeros(1, 480),
-            marks=pytest.mark.xfail,
-        ),
-        pytest.param(
+        (
             linear_gaussian,
             [
                 Gamma(ones(1), ones(1)),
@@ -257,6 +243,7 @@ def test_process_simulator(simulator: Callable, prior: Distribution):
             ],
             zeros(1, 4),
         ),
+        pytest.param(list_simulator, BoxUniform(zeros(3), ones(3)), zeros(1, 3),),
     ),
 )
 def test_prepare_sbi_problem(
@@ -273,7 +260,7 @@ def test_prepare_sbi_problem(
     simulator, prior, x_o = prepare_sbi_problem(simulator, prior, x_o)
 
     # check batch sims and type
-    n_batch = 2
+    n_batch = 1
     assert simulator(prior.sample((n_batch,))).shape[0] == n_batch
     assert isinstance(simulator(prior.sample((1,))), Tensor)
     assert prior.sample().dtype == torch.float32
@@ -294,12 +281,8 @@ def test_prepare_sbi_problem(
             np.zeros((1, 3)),
         ),
         (linear_gaussian, BoxUniform(zeros(3), ones(3)), zeros(1, 3)),
-        pytest.param(
-            list_simulator,
-            BoxUniform(zeros(3), ones(3)),
-            zeros(1, 3),
-            marks=pytest.mark.xfail,
-        ),  # list simulator.
+        (linear_gaussian_no_batch, BoxUniform(zeros(3), ones(3)), zeros(1, 3)),
+        pytest.param(list_simulator, BoxUniform(zeros(3), ones(3)), zeros(1, 3),),
         (
             numpy_linear_gaussian,
             UserNumpyUniform(zeros(3), ones(3), return_numpy=True),
@@ -333,7 +316,7 @@ def test_inference_with_user_sbi_problems(
         simulator=user_simulator,
         x_o=user_x_o,
         prior=user_prior,
-        simulation_batch_size=100,
+        simulation_batch_size=1,
     )
 
     # Run inference.