Interchangeability of Callable and BasePotential

[schroedk]

I have a question regarding the interchangeability of the argument potential_fn of

class NeuralPosterior(ABC):
    r"""Posterior $p(\theta|x)$ with `log_prob()` and `sample()` methods.<br/><br/>
    All inference methods in sbi train a neural network which is then used to obtain
    the posterior distribution. The `NeuralPosterior` class wraps the trained network
    such that one can directly evaluate the (unnormalized) log probability and draw
    samples from the posterior.
    """

    def __init__(
        self,
        potential_fn: Union[Callable, BasePotential],

For the callable case, it must be something like:

def potential(theta=None, x0=None)
    ...

in contrast to BasePotential, which is a Callable with theta as positional argument and track_gradients as keyword argument, correct? Is this tested somewhere? I only found examples where the argument is of type BasePotential.

[schroedk]

Related #1055

[janfb]

yes, correct.

The reason that the custom potential_fn has theta and x_o as args is that quantities are required to calculate the "potential", i.e., the unnormalized posterior probability.

For the BasePotential potential, the call method does not have x_o as arg, because it is set as property at runtime.

If a user passes a custom potential, then this is checked for the required args here:

sbi/sbi/inference/posteriors/base_posterior.py

Lines 57 to 69 in 593e153

	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
	)

and then wrapped as BasePotential here:

sbi/sbi/inference/potentials/base_potential.py

Lines 80 to 97 in 593e153

	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,
	prior: Optional[Distribution],
	x_o: Optional[Tensor] = None,
	device: str = "cpu",
	):
	super().__init__(prior, x_o, device)
	self.callable_potential = callable_potential
	def __call__(self, theta, track_gradients: bool = True):
	with torch.set_grad_enabled(track_gradients):
	return self.callable_potential(theta=theta, x_o=self.x_o)

[janfb]

Is this tested somewhere? I only found examples where the argument is of type BasePotential.

Yes, I had to dig a bit as well, but it's tested here:

sbi/tests/potential_test.py

Lines 28 to 37 in 593e153

	def test_callable_potential(sampling_method, mcmc_params_accurate: dict):
	"""Test whether callable potentials can be used to sample from a Gaussian."""
	dim = 2
	mean = 2.5
	cov = 2.0
	x_o = 1 * ones((dim,))
	target_density = MultivariateNormal(mean * ones((dim,)), cov * eye(dim))
	def potential(theta, x_o):
	return target_density.log_prob(theta + x_o)

Here, you can how we define a custom potential, depending on inputs theta and x_o.

[michaeldeistler]

I think this can be closed, feel free to reopen if anything is still unclear!

[janfb]

I think it's a good starting point for refactoring the Callable potential API.

[michaeldeistler]

what what have to be done here? Just more docs?

[janfb]

At the moment, if a user passes a just a Callable as potential, we test during runtime whether it has the required arguments, e.g., theta and track_gradients and x_o (or so). This is brittle. It would be nice to do this beforehand with types, e.g., define a Protocol to ensure that the passed Callable has the correct signature.