diff --git a/sbi/inference/posteriors/base_posterior.py b/sbi/inference/posteriors/base_posterior.py index 599622904..f5111cecb 100644 --- a/sbi/inference/posteriors/base_posterior.py +++ b/sbi/inference/posteriors/base_posterior.py @@ -93,7 +93,7 @@ def potential( This can be helpful for e.g. sensitivity analysis, but increases memory consumption. """ - self.potential_fn.set_x(self._x_else_default_x(x)) + self.potential_fn = self.potential_fn.bind(self._x_else_default_x(x)) theta = ensure_theta_batched(torch.as_tensor(theta)) return self.potential_fn( @@ -294,7 +294,7 @@ def map( ) if self._map is None or force_update: - self.potential_fn.set_x(self.default_x) + self.potential_fn = self.potential_fn.bind(self.default_x) self._map = self._calculate_map( num_iter=num_iter, num_to_optimize=num_to_optimize, diff --git a/sbi/inference/posteriors/ensemble_posterior.py b/sbi/inference/posteriors/ensemble_posterior.py index 755f1cebe..fa4015f92 100644 --- a/sbi/inference/posteriors/ensemble_posterior.py +++ b/sbi/inference/posteriors/ensemble_posterior.py @@ -317,7 +317,7 @@ def potential( This can be helpful for e.g. sensitivity analysis, but increases memory consumption. """ - self.potential_fn.set_x(self._x_else_default_x(x)) + self.potential_fn = self.potential_fn.bind(self._x_else_default_x(x)) theta = ensure_theta_batched(torch.as_tensor(theta)) @@ -398,7 +398,7 @@ def map( return log_weights, maps else: - self.potential_fn.set_x(self._x_else_default_x(x)) + self.potential_fn = self.potential_fn.bind(self._x_else_default_x(x)) if init_method == "posterior": inits = self.sample((num_init_samples,), self._x_else_default_x(x)) @@ -488,6 +488,28 @@ def set_x(self, x_o: Optional[Tensor]): for comp_potential in self.potential_fns: comp_potential.set_x(x_o) + def bind( + self, + x_o: Tensor, + x_is_iid: bool = True, + **kwargs, + ) -> "EnsemblePotential": + """Return a new EnsemblePotential with x_o bound to all component potentials.""" + if x_o is not None: + x_o = process_x(x_o).to(self.device) + + bound_potentials = [ + p.bind(x_o, x_is_iid=x_is_iid, **kwargs) for p in self.potential_fns + ] + + return EnsemblePotential( + potential_fns=bound_potentials, + weights=self._weights, + prior=self.prior, + x_o=x_o, + device=self.device, + ) + def __call__(self, theta: Tensor, track_gradients: bool = True) -> Tensor: r"""Returns the potential for posterior-based methods. diff --git a/sbi/inference/posteriors/importance_posterior.py b/sbi/inference/posteriors/importance_posterior.py index cfbc3d2f8..fdc98a9a1 100644 --- a/sbi/inference/posteriors/importance_posterior.py +++ b/sbi/inference/posteriors/importance_posterior.py @@ -128,7 +128,7 @@ def log_prob( `len($\theta$)`-shaped log-probability. """ x = self._x_else_default_x(x) - self.potential_fn.set_x(x) + self.potential_fn = self.potential_fn.bind(x) theta = ensure_theta_batched(torch.as_tensor(theta)) @@ -212,7 +212,7 @@ def sample( method = self.method if method is None else method - self.potential_fn.set_x(self._x_else_default_x(x)) + self.potential_fn = self.potential_fn.bind(self._x_else_default_x(x)) if method == "sir": return self._sir_sample( diff --git a/sbi/inference/posteriors/mcmc_posterior.py b/sbi/inference/posteriors/mcmc_posterior.py index f90db67d6..fc1d254cf 100644 --- a/sbi/inference/posteriors/mcmc_posterior.py +++ b/sbi/inference/posteriors/mcmc_posterior.py @@ -2,7 +2,6 @@ # under the Apache License Version 2.0, see import inspect import warnings -from copy import deepcopy from functools import partial from math import ceil from typing import Any, Callable, Dict, Literal, Optional, Union @@ -237,7 +236,7 @@ def log_prob( warn("The log-probability is unnormalized!", stacklevel=2) x = self._x_else_default_x(x) - self.potential_fn.set_x(x, x_is_iid=True) + self.potential_fn = self.potential_fn.bind(x, x_is_iid=True) theta = ensure_theta_batched(torch.as_tensor(theta)) return self.potential_fn( @@ -291,7 +290,7 @@ def sample( """ x = self._x_else_default_x(x) - self.potential_fn.set_x(x, x_is_iid=True) + self.potential_fn = self.potential_fn.bind(x, x_is_iid=True) # Replace arguments that were not passed with their default. method = self.method if method is None else method @@ -451,7 +450,7 @@ def sample_batched( # in the order of the observations. x_ = x.repeat_interleave(num_chains, dim=0) - self.potential_fn.set_x(x_, x_is_iid=False) + self.potential_fn = self.potential_fn.bind(x_, x_is_iid=False) self.potential_ = self._prepare_potential(method) # type: ignore # For each observation in the batch, we have num_chains independent chains. @@ -662,25 +661,25 @@ def _get_initial_params_batched( Tensor: initial parameters, one for each chain """ - potential_ = deepcopy(self.potential_fn) + potential_ = self.potential_fn initial_params = [] - init_fn = self._build_mcmc_init_fn( - self.proposal, - potential_fn=potential_, - transform=self.theta_transform, - init_strategy=init_strategy, # type: ignore - **kwargs, - ) for xi in x: - # Build init function - potential_.set_x(xi) + # Build init function with bound potential for this specific xi + potential_bound = potential_.bind(xi) + init_fn = self._build_mcmc_init_fn( + self.proposal, + potential_fn=potential_bound, + transform=self.theta_transform, + init_strategy=init_strategy, # type: ignore + **kwargs, + ) # Parallelize inits for resampling or sir. if num_workers > 1 and ( init_strategy == "resample" or init_strategy == "sir" ): - def seeded_init_fn(seed): + def seeded_init_fn(seed, init_fn=init_fn): torch.manual_seed(seed) return init_fn() diff --git a/sbi/inference/posteriors/rejection_posterior.py b/sbi/inference/posteriors/rejection_posterior.py index 7b05160a2..92e173129 100644 --- a/sbi/inference/posteriors/rejection_posterior.py +++ b/sbi/inference/posteriors/rejection_posterior.py @@ -120,7 +120,7 @@ def log_prob( ) warn("The log-probability is unnormalized!", stacklevel=2) - self.potential_fn.set_x(self._x_else_default_x(x)) + self.potential_fn = self.potential_fn.bind(self._x_else_default_x(x)) theta = ensure_theta_batched(torch.as_tensor(theta)) return self.potential_fn( @@ -174,7 +174,7 @@ def sample( Samples from posterior. """ num_samples = torch.Size(sample_shape).numel() - self.potential_fn.set_x(self._x_else_default_x(x)) + self.potential_fn = self.potential_fn.bind(self._x_else_default_x(x)) potential = partial(self.potential_fn, track_gradients=True) diff --git a/sbi/inference/posteriors/vector_field_posterior.py b/sbi/inference/posteriors/vector_field_posterior.py index 20c1990f0..ea27025b5 100644 --- a/sbi/inference/posteriors/vector_field_posterior.py +++ b/sbi/inference/posteriors/vector_field_posterior.py @@ -235,7 +235,7 @@ def sample( x = self._x_else_default_x(x) x = reshape_to_batch_event(x, self.vector_field_estimator.condition_shape) is_iid = x.shape[0] > 1 - self.potential_fn.set_x( + self.potential_fn = self.potential_fn.bind( x, x_is_iid=is_iid, iid_method=iid_method or self.potential_fn.iid_method, @@ -457,7 +457,9 @@ def log_prob( x = self._x_else_default_x(x) x = reshape_to_batch_event(x, self.vector_field_estimator.condition_shape) is_iid = x.shape[0] > 1 - self.potential_fn.set_x(x, x_is_iid=is_iid, **(ode_kwargs or {})) + self.potential_fn = self.potential_fn.bind( + x, x_is_iid=is_iid, **(ode_kwargs or {}) + ) theta = ensure_theta_batched(torch.as_tensor(theta)) return self.potential_fn( @@ -521,7 +523,7 @@ def sample_batched( condition_dim = len(self.vector_field_estimator.condition_shape) batch_shape = x.shape[:-condition_dim] batch_size = batch_shape.numel() - self.potential_fn.set_x(x) + self.potential_fn = self.potential_fn.bind(x) max_sampling_batch_size = ( self.max_sampling_batch_size @@ -661,7 +663,9 @@ def map( if self._map is None or force_update: # rebuild coarse flow fast for MAP optimization. - self.potential_fn.set_x(self.default_x, atol=1e-2, rtol=1e-3, exact=True) + self.potential_fn = self.potential_fn.bind( + self.default_x, atol=1e-2, rtol=1e-3, exact=True + ) callable_potential_fn = CallableDifferentiablePotentialFunction( self.potential_fn ) diff --git a/sbi/inference/posteriors/vi_posterior.py b/sbi/inference/posteriors/vi_posterior.py index e070d904f..d10162722 100644 --- a/sbi/inference/posteriors/vi_posterior.py +++ b/sbi/inference/posteriors/vi_posterior.py @@ -1229,7 +1229,7 @@ def _compute_amortized_elbo_loss(self, x_batch: Tensor, n_particles: int) -> Ten x_expanded = x_batch.repeat(n_particles, 1) # Set x_o for batched evaluation (x_is_iid=False: each θ paired with its x) - self.potential_fn.set_x(x_expanded, x_is_iid=False) + self.potential_fn = self.potential_fn.bind(x_expanded, x_is_iid=False) log_potential_flat = self.potential_fn(theta_flat) # Reshape: (n_particles * batch_size,) -> (n_particles, batch_size) diff --git a/sbi/inference/potentials/base_potential.py b/sbi/inference/potentials/base_potential.py index 8a52dd470..f353d535f 100644 --- a/sbi/inference/potentials/base_potential.py +++ b/sbi/inference/potentials/base_potential.py @@ -49,7 +49,7 @@ def x_is_iid(self) -> bool: return self._x_is_iid else: raise ValueError( - "No observed data is available. Use `potential_fn.set_x(x_o)`." + "No observed data is available. Use `potential_fn.bind(x_o)`." ) def set_x(self, x_o: Optional[Tensor], x_is_iid: Optional[bool] = True): @@ -66,9 +66,23 @@ def x_o(self) -> Tensor: return self._x_o else: raise ValueError( - "No observed data is available. Use `potential_fn.set_x(x_o)`." + "No observed data is available. Use `potential_fn.bind(x_o)`." ) + def bind(self, x_o: Tensor, x_is_iid: bool = True) -> "BasePotential": + """Create new potential with x bound, without mutable state. + + Args: + x_o: Observed data to bind. + x_is_iid: Whether x represents iid observations. + + Returns: + New potential instance with x bound. + + Subclasses must implement this method. + """ + raise NotImplementedError(f"{self.__class__.__name__} must implement bind()") + def return_x_o(self) -> Optional[Tensor]: """Return the observed data at which the potential is evaluated. @@ -149,3 +163,12 @@ def __call__(self, theta, track_gradients: bool = True): """ with torch.set_grad_enabled(track_gradients): return self.potential_fn(theta, self.x_o) + + def bind(self, x_o: Tensor, x_is_iid: bool = True) -> "CustomPotentialWrapper": + """Create new potential with x bound, without mutable state.""" + return CustomPotentialWrapper( + potential_fn=self.potential_fn, + prior=self.prior, + x_o=x_o, + device=self.device, + ) diff --git a/sbi/inference/potentials/likelihood_based_potential.py b/sbi/inference/potentials/likelihood_based_potential.py index d2907430c..991254489 100644 --- a/sbi/inference/potentials/likelihood_based_potential.py +++ b/sbi/inference/potentials/likelihood_based_potential.py @@ -94,6 +94,17 @@ def to(self, device: Union[str, torch.device]) -> "LikelihoodBasedPotential": self.likelihood_estimator.to(device) return self + def bind(self, x_o: Tensor, x_is_iid: bool = True) -> "LikelihoodBasedPotential": + """Create new potential with x bound, without mutable state.""" + bound = LikelihoodBasedPotential( + likelihood_estimator=self.likelihood_estimator, + prior=self.prior, + x_o=None, + device=self.device, + ) + bound.set_x(x_o, x_is_iid=x_is_iid) + return bound + def __call__(self, theta: Tensor, track_gradients: bool = True) -> Tensor: r"""Returns the potential $\log(p(x_o|\theta)p(\theta))$. diff --git a/sbi/inference/potentials/posterior_based_potential.py b/sbi/inference/potentials/posterior_based_potential.py index eda384626..91a4af0e3 100644 --- a/sbi/inference/potentials/posterior_based_potential.py +++ b/sbi/inference/potentials/posterior_based_potential.py @@ -100,6 +100,17 @@ def to(self, device: Union[str, torch.device]) -> "PosteriorBasedPotential": self.posterior_estimator.to(device) return self + def bind(self, x_o: Tensor, x_is_iid: bool = False) -> "PosteriorBasedPotential": + """Create new potential with x bound, without mutable state.""" + bound = PosteriorBasedPotential( + posterior_estimator=self.posterior_estimator, + prior=self.prior, + x_o=None, + device=self.device, + ) + bound.set_x(x_o, x_is_iid=x_is_iid) + 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. diff --git a/sbi/inference/potentials/ratio_based_potential.py b/sbi/inference/potentials/ratio_based_potential.py index 844405e75..917fede47 100644 --- a/sbi/inference/potentials/ratio_based_potential.py +++ b/sbi/inference/potentials/ratio_based_potential.py @@ -82,6 +82,17 @@ def to(self, device: Union[str, torch.device]) -> "RatioBasedPotential": self.ratio_estimator.to(device) return self + def bind(self, x_o: Tensor, x_is_iid: bool = True) -> "RatioBasedPotential": + """Create new potential with x bound, without mutable state.""" + bound = RatioBasedPotential( + ratio_estimator=self.ratio_estimator, + prior=self.prior, + x_o=None, + device=self.device, + ) + bound.set_x(x_o, x_is_iid=x_is_iid) + return bound + def __call__(self, theta: Tensor, track_gradients: bool = True) -> Tensor: r"""Returns the potential for likelihood-ratio-based methods. diff --git a/sbi/inference/potentials/vector_field_potential.py b/sbi/inference/potentials/vector_field_potential.py index b7ac4ba45..50681228f 100644 --- a/sbi/inference/potentials/vector_field_potential.py +++ b/sbi/inference/potentials/vector_field_potential.py @@ -61,6 +61,7 @@ def __init__( self.vector_field_estimator.eval() self.iid_method = iid_method self.iid_params = iid_params + self.neural_ode_backend = neural_ode_backend neural_ode_kwargs = neural_ode_kwargs or {} self.neural_ode = build_neural_ode( @@ -128,6 +129,40 @@ def set_x( elif self._x_o is not None: self.flows = self.rebuild_flows_for_batch(**ode_kwargs) + def bind( + self, + x_o: Tensor, + x_is_iid: bool = False, + iid_method: Optional[str] = None, + iid_params: Optional[Dict[str, Any]] = None, + guidance_method: Optional[str] = None, + guidance_params: Optional[Dict[str, Any]] = None, + **ode_kwargs, + ) -> "VectorFieldBasedPotential": + """Create new potential with x bound, without mutable state.""" + bound = VectorFieldBasedPotential( + vector_field_estimator=self.vector_field_estimator, + prior=self.prior, + x_o=None, + device=self.device, + iid_method=self.iid_method, + iid_params=self.iid_params, + neural_ode_backend=self.neural_ode_backend, + ) + bound.neural_ode.params.update(self.neural_ode.params) + # Transfer learned neural ODE parameters (e.g., from training) to the bound + # potential. This preserves the trained flow model when conditioning on new x. + bound.set_x( + 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, + ) + return bound + def __call__( self, theta: Tensor, diff --git a/sbi/utils/conditional_density_utils.py b/sbi/utils/conditional_density_utils.py index 3251f99a9..499981b8c 100644 --- a/sbi/utils/conditional_density_utils.py +++ b/sbi/utils/conditional_density_utils.py @@ -408,11 +408,13 @@ def __call__(self, theta: Tensor, track_gradients: bool = True) -> Tensor: def x_is_iid(self) -> bool: """If x has batch dimension greater than 1, whether to intepret the batch as iid samples or batch of data points.""" - if self._x_is_iid is not None: + if hasattr(self, "_x_is_iid") and self._x_is_iid is not None: return self._x_is_iid + elif hasattr(self, "potential_fn") and self.potential_fn is not None: + return self.potential_fn.x_is_iid else: raise ValueError( - "No observed data is available. Use `potential_fn.set_x(x_o)`." + "No observed data is available. Use `potential_fn.bind(x_o)`." ) def set_x(self, x_o: Optional[Tensor], x_is_iid: Optional[bool] = True): @@ -420,7 +422,17 @@ def set_x(self, x_o: Optional[Tensor], x_is_iid: Optional[bool] = True): 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.set_x(x_o, x_is_iid=x_is_iid) + self.potential_fn = self.potential_fn.bind(x_o, x_is_iid=x_is_iid) + + def bind(self, x_o: Tensor, x_is_iid: bool = True) -> "ConditionedPotential": + """Create new potential with x bound, without mutable state.""" + bound = ConditionedPotential( + potential_fn=self.potential_fn.bind(x_o, x_is_iid), + condition=self.condition, + dims_to_sample=self.dims_to_sample, + ) + bound._x_is_iid = x_is_iid + return bound @property def x_o(self) -> Tensor: diff --git a/tests/inference_on_device_test.py b/tests/inference_on_device_test.py index 395f09a71..a56f323fa 100644 --- a/tests/inference_on_device_test.py +++ b/tests/inference_on_device_test.py @@ -396,6 +396,12 @@ def __call__(self, theta, **kwargs): def allow_iid_x(self) -> bool: return True + 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) + return bound + potential_fn = FakePotential( prior=MultivariateNormal( zeros(num_dim, device=device), eye(num_dim, device=device) @@ -450,6 +456,12 @@ def __call__(self, theta, **kwargs): def allow_iid_x(self) -> bool: return True + 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) + return bound + potential_fn = FakePotential(prior=prior, device=device) # Generate mock simulation data on device diff --git a/tests/vi_test.py b/tests/vi_test.py index e1327e56a..f3f187ce9 100644 --- a/tests/vi_test.py +++ b/tests/vi_test.py @@ -57,6 +57,12 @@ def __call__(self, theta, **kwargs): def allow_iid_x(self) -> bool: return True + 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) + return bound + def make_tractable_potential(target_distribution, prior): """Create a potential function from a known target distribution.""" @@ -70,6 +76,14 @@ def __call__(self, theta, **kwargs): def allow_iid_x(self) -> bool: return True + def bind( + self, x_o: torch.Tensor, x_is_iid: bool = True + ) -> "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) + return bound + return TractablePotential(prior=prior)