Type 'potential_fn' properly, not 'Callable'

[Baschdl]

We currently type potential_fn only as Callable

sbi/sbi/utils/conditional_density_utils.py

Lines 273 to 276 in bae6994

	class ConditionedPotential:
	def __init__(
	self,
	potential_fn: Callable,

despite implicitly requiring it to be a BasePotential to be able to call something like set_x:

sbi/sbi/utils/conditional_density_utils.py

Line 330 in bae6994

self.potential_fn.set_x(x_o)

Changing the type to BasePotential will lead to other required changes down the line.

[janfb]

Note that we have a "hacky" runtime check for this implemented in the base_posterior.py (but not in ConditionedPotential).

sbi/sbi/inference/posteriors/base_posterior.py

Lines 55 to 69 in 366b67f

	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
	potential_fn = CallablePotentialWrapper(
	potential_fn, prior=None, x_o=None, device=potential_device
	)

From a discussion with my colleagues at Transferlab, I would suggest using a Protocol to define the required properties of the Callable. Here is an outline for implementing this:

Step 1: Define a Protocol

Define a Protocol that includes a method with the required arguments (theta and x_o in this case). This method will define the expected signature for potential_fn.

from typing import Protocol, Any

class PotentialProtocol(Protocol):
    def __call__(self, *, theta: Any, x_o: Any) -> float:
        ...

Step 2: Type Checking with the Protocol

Instead of checking the arguments at runtime, you use the PotentialProtocol as a type hint. This way, any function or callable object passed as potential_fn must adhere to this protocol.

Example Usage in Class

from typing import Any
import inspect

class ConditionedPotential:
    def __init__(self, potential_fn: PotentialProtocol) -> None:
        self.potential_fn = potential_fn

        # Optional: Runtime check if not using static type checking tools
        if not isinstance(potential_fn, BasePotential):
            self._validate_callable_signature(potential_fn)

    def _validate_callable_signature(self, fn: PotentialProtocol) -> None:
        # Get the signature of the callable
        kwargs_of_callable = list(inspect.signature(fn).parameters.keys())
        
        # Ensure required arguments are present
        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."
            )

    def some_method(self, theta: Any, x_o: Any) -> float:
        return self.potential_fn(theta=theta, x_o=x_o)

Explanation

1. PotentialProtocol Definition:

   • This protocol defines a method __call__ with keyword arguments theta and x_o. The method returns a float. This setup ensures that any callable passed as potential_fn must implement this signature.

2. ConditionedPotential Class:

   • The constructor now uses PotentialProtocol to type the potential_fn argument. This provides a clear expectation of the function's interface.
   • The _validate_callable_signature method is an optional runtime check to ensure that the callable conforms to the expected signature, useful if you're not relying solely on static type checking.

3. Static Type Checking:

   • Using tools like mypy, you can catch potential issues at development time, ensuring that all functions passed as potential_fn conform to the PotentialProtocol.

Additional Considerations

• BasePotential: If you have existing classes that inherit from BasePotential, ensure they conform to the PotentialProtocol either directly or by implementing the necessary methods.

• Runtime vs. Static Checking: The static type checking provided by the protocol can catch issues during development. The runtime check (_validate_callable_signature) can be kept for additional safety, especially if you're unsure about the inputs or when integrating with dynamically typed parts of the code.

This approach formalizes the expected interface and improves both clarity and safety, reducing the likelihood of runtime errors due to incorrectly structured callables.

[schroedk]

@janfb as discussed last week, the argument potential_fn of class ConditionalPotential must be of type BasePotential.
I created another issue, regarding our discussion of the different callable types: #1223