diff --git a/sbi/inference/posteriors/ensemble_posterior.py b/sbi/inference/posteriors/ensemble_posterior.py index fa4015f92..042c19265 100644 --- a/sbi/inference/posteriors/ensemble_posterior.py +++ b/sbi/inference/posteriors/ensemble_posterior.py @@ -1,6 +1,7 @@ # This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed # under the Apache License Version 2.0, see +import warnings from typing import List, Optional, Tuple, Union import torch @@ -479,14 +480,21 @@ def allow_iid_x(self) -> bool: ) def set_x(self, x_o: Optional[Tensor]): - """Check the shape of the observed data and, if valid, set it.""" - if x_o is not None: - x_o = process_x(x_o).to( # type: ignore - self.device - ) - self._x_o = x_o - for comp_potential in self.potential_fns: - comp_potential.set_x(x_o) + """Check the shape of the observed data and, if valid, set it. + + DEPRECATED: Use bind() instead. This method delegates to bind() internally. + It will be removed in a future release. + """ + + warnings.warn( + "set_x() is deprecated and will be removed in a future release. " + "Use bind() instead.", + FutureWarning, + stacklevel=2, + ) + bound = self.bind(x_o) + self._x_o = bound._x_o + self.potential_fns = bound.potential_fns def bind( self, diff --git a/sbi/inference/potentials/base_potential.py b/sbi/inference/potentials/base_potential.py index 7b9676812..946e9b134 100644 --- a/sbi/inference/potentials/base_potential.py +++ b/sbi/inference/potentials/base_potential.py @@ -56,11 +56,20 @@ def x_is_iid(self) -> bool: ) def set_x(self, x_o: Optional[Tensor], x_is_iid: Optional[bool] = True): - """Check the shape of the observed data and, if valid, set it.""" - if x_o is not None: - x_o = process_x(x_o).to(self.device) - self._x_o = x_o - self._x_is_iid = x_is_iid + """Check the shape of the observed data and, if valid, set it. + + DEPRECATED: Use bind() instead. This method delegates to bind() internally. + """ + import warnings + + warnings.warn( + "set_x() is deprecated, use bind() instead", + FutureWarning, + stacklevel=2, + ) + bound = self.bind(x_o, x_is_iid=x_is_iid) + self._x_o = bound._x_o + self._x_is_iid = bound._x_is_iid @property def x_o(self) -> Tensor: diff --git a/sbi/inference/potentials/posterior_based_potential.py b/sbi/inference/potentials/posterior_based_potential.py index 791bc8c30..4c82d0cc8 100644 --- a/sbi/inference/potentials/posterior_based_potential.py +++ b/sbi/inference/potentials/posterior_based_potential.py @@ -115,13 +115,22 @@ def bind(self, x_o: Tensor, x_is_iid: bool = False) -> "PosteriorBasedPotential" return bound def set_x(self, x_o: Optional[Tensor], x_is_iid: Optional[bool] = False): + """Check the shape of the observed data and, if valid, set it. + + DEPRECATED: Use bind() instead. This method delegates to bind() internally. + It will be removed in a future release. """ - Check the shape of the observed data and, if valid, set it. - """ - if x_o is not None: - x_o = process_x(x_o).to(self.device) - self._x_o = x_o - self._x_is_iid = x_is_iid + import warnings + + warnings.warn( + "set_x() is deprecated and will be removed in a future release. " + "Use bind() instead.", + FutureWarning, + stacklevel=2, + ) + bound = self.bind(x_o, x_is_iid=x_is_iid) + self._x_o = bound._x_o + self._x_is_iid = bound._x_is_iid def __call__(self, theta: Tensor, track_gradients: bool = True) -> Tensor: r"""Returns the potential for posterior-based methods. diff --git a/sbi/inference/potentials/vector_field_potential.py b/sbi/inference/potentials/vector_field_potential.py index 8c492eb53..7ce061e49 100644 --- a/sbi/inference/potentials/vector_field_potential.py +++ b/sbi/inference/potentials/vector_field_potential.py @@ -1,6 +1,7 @@ # This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed # under the Apache License Version 2.0, see +import warnings from typing import Any, Dict, List, Literal, Optional, Tuple, Union import torch @@ -113,6 +114,8 @@ def set_x( Rebuilds the continuous normalizing flow if the observed data is set. + DEPRECATED: Use bind() instead. This method delegates to bind() internally. + It will be removed in a future release. Args: x_o: The observed data. x_is_iid: Whether the observed data is IID (if batch_dim>1). @@ -122,18 +125,32 @@ def set_x( `IIDScoreFunction`. ode_kwargs: Additional keyword arguments for the neural ODE. """ - if x_o is not None: - x_o = process_x(x_o).to(self.device) - self._x_o = x_o - self._x_is_iid = x_is_iid - self.iid_method = iid_method or self.iid_method - self.iid_params = iid_params - self.guidance_method = guidance_method - self.guidance_params = guidance_params - if not x_is_iid and (self._x_o is not None): - self.flow = self.rebuild_flow(**ode_kwargs) - elif self._x_o is not None: - self.flows = self.rebuild_flows_for_batch(**ode_kwargs) + + warnings.warn( + "set_x() is deprecated and will be removed in a future release. " + "Use bind() instead.", + FutureWarning, + stacklevel=2, + ) + bound = self.bind( + x_o, + x_is_iid=x_is_iid, + iid_method=iid_method, + iid_params=iid_params, + guidance_method=guidance_method, + guidance_params=guidance_params, + **ode_kwargs, + ) + self._x_o = bound._x_o + self._x_is_iid = bound._x_is_iid + self.iid_method = bound.iid_method + self.iid_params = bound.iid_params + self.guidance_method = bound.guidance_method + self.guidance_params = bound.guidance_params + if not x_is_iid and (bound._x_o is not None): + self.flow = bound.flow + elif bound._x_o is not None: + self.flows = bound.flows def bind( self, diff --git a/sbi/utils/conditional_density_utils.py b/sbi/utils/conditional_density_utils.py index 499981b8c..bb2bfdd1f 100644 --- a/sbi/utils/conditional_density_utils.py +++ b/sbi/utils/conditional_density_utils.py @@ -15,7 +15,6 @@ MixtureDensityEstimator, ) from sbi.utils.torchutils import ensure_theta_batched -from sbi.utils.user_input_checks import process_x def compute_corrcoeff(probs: Tensor, limits: Tensor): @@ -418,11 +417,20 @@ def x_is_iid(self) -> bool: ) def set_x(self, x_o: Optional[Tensor], x_is_iid: Optional[bool] = True): - """Check the shape of the observed data and, if valid, set it.""" - if x_o is not None: - x_o = process_x(x_o).to(self.device) - self._x_is_iid = x_is_iid - self.potential_fn = self.potential_fn.bind(x_o, x_is_iid=x_is_iid) + """Check the shape of the observed data and, if valid, set it. + + DEPRECATED: Use bind() instead. This method delegates to bind() internally. + """ + import warnings + + warnings.warn( + "set_x() is deprecated, use bind() instead", + FutureWarning, + stacklevel=2, + ) + bound = self.bind(x_o, x_is_iid=x_is_iid) + self._x_is_iid = bound._x_is_iid + self.potential_fn = bound.potential_fn def bind(self, x_o: Tensor, x_is_iid: bool = True) -> "ConditionedPotential": """Create new potential with x bound, without mutable state.""" diff --git a/tests/inference_on_device_test.py b/tests/inference_on_device_test.py index a56f323fa..53d927f8f 100644 --- a/tests/inference_on_device_test.py +++ b/tests/inference_on_device_test.py @@ -51,7 +51,7 @@ from sbi.utils import BoxUniform from sbi.utils.sbiutils import seed_all_backends from sbi.utils.torchutils import gpu_available, process_device -from sbi.utils.user_input_checks import validate_theta_and_x +from sbi.utils.user_input_checks import process_x, validate_theta_and_x pytestmark = pytest.mark.skipif( not gpu_available(), reason="No CUDA or MPS device available." @@ -399,7 +399,9 @@ def allow_iid_x(self) -> bool: def bind(self, x_o: torch.Tensor, x_is_iid: bool = True) -> "FakePotential": """Create new potential with x bound, without mutable state.""" bound = FakePotential(prior=self.prior, device=self.device) - bound.set_x(x_o, x_is_iid=x_is_iid) + x_o = process_x(x_o).to(self.device) + bound._x_o = x_o + bound._x_is_iid = x_is_iid return bound potential_fn = FakePotential( @@ -459,7 +461,9 @@ def allow_iid_x(self) -> bool: def bind(self, x_o: torch.Tensor, x_is_iid: bool = True) -> "FakePotential": """Create new potential with x bound, without mutable state.""" bound = FakePotential(prior=self.prior, device=self.device) - bound.set_x(x_o, x_is_iid=x_is_iid) + x_o = process_x(x_o).to(self.device) + bound._x_o = x_o + bound._x_is_iid = x_is_iid return bound potential_fn = FakePotential(prior=prior, device=device) diff --git a/tests/vi_test.py b/tests/vi_test.py index f3f187ce9..9aac865cf 100644 --- a/tests/vi_test.py +++ b/tests/vi_test.py @@ -34,6 +34,7 @@ ) from sbi.utils import MultipleIndependent from sbi.utils.metrics import c2st, check_c2st +from sbi.utils.user_input_checks import process_x # Supported variational families for VI FLOWS = ["maf", "nsf", "naf", "unaf", "nice", "sospf", "gaussian", "gaussian_diag"] @@ -60,7 +61,9 @@ def allow_iid_x(self) -> bool: def bind(self, x_o: torch.Tensor, x_is_iid: bool = True) -> "FakePotential": """Create new potential with x bound, without mutable state.""" bound = FakePotential(prior=self.prior, device=self.device) - bound.set_x(x_o, x_is_iid=x_is_iid) + x_o = process_x(x_o).to(self.device) + bound._x_o = x_o + bound._x_is_iid = x_is_iid return bound @@ -81,7 +84,9 @@ def bind( ) -> "TractablePotential": """Create new potential with x bound, without mutable state.""" bound = TractablePotential(prior=self.prior, device=self.device) - bound.set_x(x_o, x_is_iid=x_is_iid) + x_o = process_x(x_o).to(self.device) + bound._x_o = x_o + bound._x_is_iid = x_is_iid return bound return TractablePotential(prior=prior)