Slight improvements

Author: manuelgloeckler

sbi/inference/posteriors/vector_field_posterior.py

@@ -407,10 +407,7 @@ def log_prob(
         x = self._x_else_default_x(x)
         x = reshape_to_batch_event(x, self.vector_field_estimator.condition_shape)
         is_iid = x.shape[0] > 1
-        self.potential_fn.set_x(
-            x,
-            x_is_iid=is_iid,
-        )
+        self.potential_fn.set_x(x, x_is_iid=is_iid, **(ode_kwargs or {}))
 
         theta = ensure_theta_batched(torch.as_tensor(theta))
         return self.potential_fn(

sbi/inference/potentials/vector_field_potential.py

@@ -150,7 +150,7 @@ def __call__(
                 assert self.flows is not None, (
                     "Flows for each iid x are required for evaluating log_prob."
                 )
-                n = self.x_o.shape[0]  # number of iid samples
+                num_iid = self.x_o.shape[0]  # number of iid samples
                 iid_posteriors_prob = torch.sum(
                     torch.stack(
                         [
@@ -162,8 +162,8 @@ def __call__(
                     dim=0,
                 )
                 # Apply the adjustment for iid observations i.e. we have to subtract
-                # (n-1) times the log prior.
-                log_probs = iid_posteriors_prob - (n - 1) * self.prior.log_prob(
+                # (num_iid-1) times the log prior.
+                log_probs = iid_posteriors_prob - (num_iid - 1) * self.prior.log_prob(
                     theta_density_estimator
                 ).squeeze(-1)
             else:
@@ -243,8 +243,8 @@ def rebuild_flow(self, **kwargs) -> NormalizingFlow:
         """
         if self._x_o is None:
             raise ValueError(
-                "No observed data x_o is available. Please reinitialize \
-                the potential or manually set self._x_o."
+                "No observed data x_o is available. Please reinitialize"
+                "the potential or manually set self._x_o."
             )
         x_density_estimator = reshape_to_batch_event(
             self.x_o, event_shape=self.vector_field_estimator.condition_shape
@@ -253,17 +253,15 @@ def rebuild_flow(self, **kwargs) -> NormalizingFlow:
         flow = self.neural_ode(x_density_estimator, **kwargs)
         return flow
 
-    def rebuild_flows_for_batch(
-        self, atol: float = 1e-5, rtol: float = 1e-6, exact: bool = True
-    ) -> List[NormalizingFlow]:
+    def rebuild_flows_for_batch(self, **kwargs) -> List[NormalizingFlow]:
         """
         Rebuilds the continuous normalizing flows for each iid in x_o. This is used when
         a new default x_o is set, or to evaluate the log probs at higher precision.
         """
         if self._x_o is None:
             raise ValueError(
-                "No observed data x_o is available. Please reinitialize \
-                the potential or manually set self._x_o."
+                "No observed data x_o is available. Please reinitialize "
+                "the potential or manually set self._x_o."
             )
         flows = []
         for i in range(self._x_o.shape[0]):
@@ -272,9 +270,7 @@ def rebuild_flows_for_batch(
                 iid_x, event_shape=self.vector_field_estimator.condition_shape
             )
 
-            flow = self.neural_ode(
-                condition=x_density_estimator, atol=atol, rtol=rtol, exact=exact
-            )
+            flow = self.neural_ode(x_density_estimator, **kwargs)
             flows.append(flow)
         return flows
 

tests/linearGaussian_vector_field_test.py

@@ -593,7 +593,7 @@ def simulator(theta):
 @pytest.mark.slow
 @pytest.mark.parametrize("vector_field_type", ["ve", "vp", "fmpe"])
 @pytest.mark.parametrize("prior_type", ["gaussian"])
-@pytest.mark.parametrize("iid_batch_size", [1, 2])
+@pytest.mark.parametrize("iid_batch_size", [1, 2, 5])
 def test_iid_log_prob(vector_field_type, prior_type, iid_batch_size):
     '''
     Tests the log-probability computation of the score-based posterior.
@@ -630,7 +630,7 @@ def test_iid_log_prob(vector_field_type, prior_type, iid_batch_size):
     approx_prob = approx_posterior.log_prob(posterior_samples, x=x_o)
 
     diff = torch.abs(true_prob - approx_prob)
-    assert diff.mean() < 0.4, (
+    assert diff.mean() < 0.3 * iid_batch_size, (
         f"Probs diff: {diff.mean()} too big "
         f"for number of samples {num_posterior_samples}"
     )