diff --git a/sbi/inference/posteriors/base_posterior.py b/sbi/inference/posteriors/base_posterior.py index 599622904..ab291994a 100644 --- a/sbi/inference/posteriors/base_posterior.py +++ b/sbi/inference/posteriors/base_posterior.py @@ -79,6 +79,14 @@ def __init__( # as default x. self._x = self.potential_fn.return_x_o() + self._initialized = False + + def init(self): + """Initialize the underlying potential before sampling.""" + self.potential_fn.init() + self._initialized = True + return self + def potential( self, theta: Tensor, x: Optional[Tensor] = None, track_gradients: bool = False ) -> Tensor: @@ -93,6 +101,8 @@ def potential( This can be helpful for e.g. sensitivity analysis, but increases memory consumption. """ + if not self._initialized: + self.init() self.potential_fn.set_x(self._x_else_default_x(x)) theta = ensure_theta_batched(torch.as_tensor(theta)) diff --git a/sbi/inference/potentials/base_potential.py b/sbi/inference/potentials/base_potential.py index 04e62ef09..6bb30fec6 100644 --- a/sbi/inference/potentials/base_potential.py +++ b/sbi/inference/potentials/base_potential.py @@ -30,6 +30,12 @@ def __init__( self.device = device self.prior = prior self.set_x(x_o) + self._initialized = False + + def init(self): + """Initialize potential (can be overridden by subclasses).""" + self._initialized = True + return self @abstractmethod def __call__(self, theta: Tensor, track_gradients: bool = True) -> Tensor: @@ -136,5 +142,7 @@ def __call__(self, theta, track_gradients: bool = True): Note, x_o is re-used from the initialization of the potential function. """ + if not self._initialized: + self.init() with torch.set_grad_enabled(track_gradients): return self.potential_fn(theta, self.x_o)