refactor: callable potential fun; add test

Author: janfb

sbi/inference/posteriors/base_posterior.py

@@ -1,7 +1,6 @@
 # This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
 # under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
 
-import inspect
 from abc import abstractmethod
 from typing import Any, Callable, Dict, Optional, Union
 from warnings import warn
@@ -54,14 +53,6 @@ def __init__(
 
         # Wrap as `CallablePotentialWrapper` if `potential_fn` is a Callable.
         if not isinstance(potential_fn, BasePotential):
-            kwargs_of_callable = list(inspect.signature(potential_fn).parameters.keys())
-            for key in ["theta", "x_o"]:
-                assert key in kwargs_of_callable, (
-                    "If you pass a `Callable` as `potential_fn` then it must have "
-                    "`theta` and `x_o` as inputs, even if some of these keyword "
-                    "arguments are unused."
-                )
-
             # If the `potential_fn` is a Callable then we wrap it as a
             # `CallablePotentialWrapper` which inherits from `BasePotential`.
             potential_device = "cpu" if device is None else device

sbi/inference/potentials/base_potential.py

@@ -1,8 +1,9 @@
 # This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
 # under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
 
+import inspect
 from abc import ABCMeta, abstractmethod
-from typing import Optional
+from typing import Callable, Optional
 
 import torch
 from torch import Tensor
@@ -85,18 +86,39 @@ def return_x_o(self) -> Optional[Tensor]:
 class CallablePotentialWrapper(BasePotential):
     """If `potential_fn` is a callable it gets wrapped as this."""
 
-    allow_iid_x = True  # type: ignore
-
     def __init__(
         self,
-        callable_potential,
+        potential_fn: Callable,
         prior: Optional[Distribution],
         x_o: Optional[Tensor] = None,
         device: str = "cpu",
     ):
+        """Wraps a callable potential function.
+
+        Args:
+            potential_fn: Callable potential function, must have `theta` and `x_o` as
+                arguments.
+            prior: Prior distribution.
+            x_o: Observed data.
+            device: Device on which to evaluate the potential function.
+
+        """
         super().__init__(prior, x_o, device)
-        self.callable_potential = callable_potential
+
+        kwargs_of_callable = list(inspect.signature(potential_fn).parameters.keys())
+        required_keys = ["theta", "x_o"]
+        for key in required_keys:
+            assert key in kwargs_of_callable, (
+                "If you pass a `Callable` as `potential_fn` then it must have "
+                "`theta` and `x_o` as inputs, even if some of these keyword "
+                "arguments are unused."
+            )
+        self.potential_fn = potential_fn
 
     def __call__(self, theta, track_gradients: bool = True):
+        """Call the callable potential function on given theta.
+
+        Note, x_o is re-used from the initialization of the potential function.
+        """
         with torch.set_grad_enabled(track_gradients):
-            return self.callable_potential(theta=theta, x_o=self.x_o)
+            return self.potential_fn(theta=theta, x_o=self.x_o)

tests/potential_test.py

@@ -5,7 +5,7 @@
 
 import pytest
 import torch
-from torch import eye, ones, zeros
+from torch import Tensor, eye, ones, zeros
 from torch.distributions import MultivariateNormal
 
 from sbi.inference import (
@@ -14,6 +14,9 @@
     RejectionPosterior,
     VIPosterior,
 )
+from sbi.inference.potentials.base_potential import CallablePotentialWrapper
+from sbi.utils import BoxUniform
+from sbi.utils.conditional_density_utils import ConditionedPotential
 
 
 @pytest.mark.parametrize(
@@ -64,3 +67,25 @@ def potential(theta, x_o):
     sample_std = torch.std(approx_samples, dim=0)
     assert torch.allclose(sample_mean, torch.as_tensor(mean) - x_o, atol=0.2)
     assert torch.allclose(sample_std, torch.sqrt(torch.as_tensor(cov)), atol=0.1)
+
+
+@pytest.mark.parametrize(
+    "condition",
+    [
+        torch.rand(1, 2),
+        pytest.param(
+            torch.rand(2, 2),
+            marks=pytest.mark.xfail(
+                raises=ValueError,
+                match="Condition with batch size > 1 not supported",
+            ),
+        ),
+    ],
+)
+def test_conditioned_potential(condition: Tensor):
+    potential_fn = CallablePotentialWrapper(
+        potential_fn=lambda theta, x_o: theta,
+        prior=BoxUniform(low=zeros(2), high=ones(2)),
+    )
+
+    ConditionedPotential(potential_fn, condition=condition, dims_to_sample=[0])