fix: add protocol for custom potential (#1409)

* fix: protocol and refactor for custom potential

* fix docstrings

* fix docstring

Author: janfb

sbi/inference/posteriors/base_posterior.py

@@ -2,7 +2,7 @@
 # under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
 
 from abc import abstractmethod
-from typing import Any, Callable, Dict, Optional, Union
+from typing import Any, Dict, Optional, Union
 from warnings import warn
 
 import torch
@@ -11,7 +11,8 @@
 
 from sbi.inference.potentials.base_potential import (
     BasePotential,
-    CallablePotentialWrapper,
+    CustomPotential,
+    CustomPotentialWrapper,
 )
 from sbi.sbi_types import Array, Shape, TorchTransform
 from sbi.utils.sbiutils import gradient_ascent
@@ -29,7 +30,7 @@ class NeuralPosterior:
 
     def __init__(
         self,
-        potential_fn: Union[Callable, BasePotential],
+        potential_fn: Union[BasePotential, CustomPotential],
         theta_transform: Optional[TorchTransform] = None,
         device: Optional[str] = None,
         x_shape: Optional[torch.Size] = None,
@@ -51,16 +52,13 @@ def __init__(
                 stacklevel=2,
             )
 
-        # Wrap as `CallablePotentialWrapper` if `potential_fn` is a Callable.
+        # Wrap custom potential functions to adhere to the `BasePotential` interface.
         if not isinstance(potential_fn, BasePotential):
-            # 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
-            potential_fn = CallablePotentialWrapper(
+            potential_fn = CustomPotentialWrapper(
                 potential_fn, prior=None, x_o=None, device=potential_device
             )
 
-        # Ensure device string.
         self._device = process_device(potential_fn.device if device is None else device)
 
         self.potential_fn = potential_fn

sbi/inference/potentials/base_potential.py

@@ -1,9 +1,8 @@
 # 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 Callable, Optional
+from typing import Optional, Protocol
 
 import torch
 from torch import Tensor
@@ -83,42 +82,43 @@ def return_x_o(self) -> Optional[Tensor]:
         return self._x_o
 
 
-class CallablePotentialWrapper(BasePotential):
+class CustomPotential(Protocol):
+    """Protocol for custom potential functions."""
+
+    def __call__(self, theta: Tensor, x_o: Tensor) -> Tensor:
+        """Call the potential function on given theta and observed data."""
+        ...
+
+
+class CustomPotentialWrapper(BasePotential):
     """If `potential_fn` is a callable it gets wrapped as this."""
 
     def __init__(
         self,
-        potential_fn: Callable,
+        potential_fn: CustomPotential,
         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.
+            potential_fn: Custom potential function following the CustomPotential
+                protocol, i.e., the function must have exactly two positional arguments
+                where the first is theta and the second is the x_o.
+            prior: Prior distribution, optional at init, but needed at inference time.
+            x_o: Observed data, optional at init, but needed at inference time.
             device: Device on which to evaluate the potential function.
 
         """
         super().__init__(prior, x_o, device)
 
-        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.
+        """Calls the custom 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.potential_fn(theta=theta, x_o=self.x_o)
+            return self.potential_fn(theta, self.x_o)

tests/potential_test.py

@@ -14,7 +14,7 @@
     RejectionPosterior,
     VIPosterior,
 )
-from sbi.inference.potentials.base_potential import CallablePotentialWrapper
+from sbi.inference.potentials.base_potential import CustomPotentialWrapper
 from sbi.utils import BoxUniform
 from sbi.utils.conditional_density_utils import ConditionedPotential
 
@@ -83,7 +83,7 @@ def potential(theta, x_o):
     ],
 )
 def test_conditioned_potential(condition: Tensor):
-    potential_fn = CallablePotentialWrapper(
+    potential_fn = CustomPotentialWrapper(
         potential_fn=lambda theta, x_o: theta,
         prior=BoxUniform(low=zeros(2), high=ones(2)),
     )