def _expected_value_f_of_x(probs: Tensor, limits: Tensor, f: Callable = lambda x: x) -> Tensor: """ Return the expected value of a function of random variable(s) E[f(X_i,...,X_k)]. The expected value is computed from evaluations of the joint density on an evenly spaced grid, passed as `probs`. This function can not deal with functions `f` that have multiple outputs. They will simply be summed over. Args: probs: Matrix of evaluations of the density. limits: Limits within which the entries of the matrix are evenly spaced. f: The operation to be applied to the expected values. Returns: Expected value. """ probs = ensure_theta_batched(probs) limits = ensure_theta_batched(limits) x_values_over_which_we_integrate = [ torch.linspace(lim[0], lim[1], prob.shape[0]) for lim, prob in zip(torch.flip(limits, [0]), probs) ] # See #403 and #404 for flip(). grids = list(torch.meshgrid(x_values_over_which_we_integrate)) expected_val = torch.sum(f(*grids) * probs) limits_diff = torch.prod(limits[:, 1] - limits[:, 0]) expected_val /= probs.numel() / limits_diff.item() return expected_val
def _prepare_theta_and_x_for_log_prob_( self, theta: Tensor, x: Optional[Tensor] = None) -> Tuple[Tensor, Tensor]: r"""Returns $\theta$ and $x$ in shape that can be used by posterior.log_prob(). Checks shapes of $\theta$ and $x$ and then repeats $x$ as often as there were batch elements in $\theta$. Moves $\theta$ and $x$ to the device of the neural net. Args: theta: Parameters $\theta$. x: Conditioning context for posterior $p(\theta|x)$. If not provided, fall back onto an `x_o` if previously provided for multi-round training, or to another default if set later for convenience, see `.set_default_x()`. Returns: ($\theta$, $x$) with the same batch dimension, where $x$ is repeated as often as there were batch elements in $\theta$ originally. """ theta = ensure_theta_batched(torch.as_tensor(theta)) # Select and check x to condition on. x = atleast_2d_float32_tensor(self._x_else_default_x(x)) if not self._allow_iid_x: self._ensure_single_x(x) self._ensure_x_consistent_with_default_x(x) return theta, x
def np_potential(self, theta: np.array) -> ScalarFloat: """Return potential for Numpy slice sampler." Args: theta: Parameters $\theta$, batch dimension 1. Returns: Posterior log probability of theta. """ theta = torch.as_tensor(theta, dtype=torch.float32) theta = ensure_theta_batched(theta) num_batch = theta.shape[0] x_batched = ensure_x_batched(self.x) # Repeat x over batch dim to match theta batch, accounting for multi-D x. x_repeated = x_batched.repeat(num_batch, *(1 for _ in range(x_batched.ndim - 1))) assert ( x_batched.ndim == 2 ), """X must not be multidimensional for ratio-based methods because it will be concatenated with theta.""" with torch.set_grad_enabled(False): log_ratio = (self.classifier( torch.cat((theta.to(self.x.device), x_repeated), dim=1)).reshape(-1).cpu()) # Notice opposite sign to pyro potential. return log_ratio + self.prior.log_prob(theta)
def log_prob( self, theta: Tensor, x: Optional[Tensor] = None, track_gradients: bool = False ) -> Tensor: r"""Returns the log-probability of theta under the posterior. Args: theta: Parameters $\theta$. track_gradients: Whether the returned tensor supports tracking gradients. This can be helpful for e.g. sensitivity analysis, but increases memory consumption. Returns: `len($\theta$)`-shaped log-probability. """ warn( "`.log_prob()` is deprecated for methods that can only evaluate the log-probability up to a normalizing constant. Use `.potential()` instead." ) warn("The log-probability is unnormalized!") self.potential_fn.set_x(self._x_else_default_x(x)) theta = ensure_theta_batched(torch.as_tensor(theta)) return self.potential_fn( theta.to(self._device), track_gradients=track_gradients )
def __call__(self, theta: Tensor, track_gradients: bool = True) -> Tensor: r""" Returns the conditional potential $\log(p(\theta_i|\theta_j, x))$. Args: theta: Free parameters $\theta_i$, batch dimension 1. Returns: Conditional posterior log-probability $\log(p(\theta_i|\theta_j, x))$, masked outside of prior. """ theta_ = ensure_theta_batched( torch.as_tensor(theta, dtype=torch.float32)) # `theta_condition`` will first have all entries of the `condition` and then # override the entries that should be sampled with `theta` (see below). theta_condition = deepcopy(self.condition) # In case `theta` is a batch of theta (e.g. multi-chain MCMC), we have to # repeat `theta_condition`` to the same batchsize. theta_condition = theta_condition.repeat(theta_.shape[0], 1) theta_condition[:, self.dims_to_sample] = theta_ return self.potential_fn(theta_condition, track_gradients=track_gradients)
def _prepare_theta_and_x_for_log_prob_( self, theta: Tensor, x: Optional[Tensor] = None, ) -> Tuple[Tensor, Tensor]: r"""Returns $\theta$ and $x$ in shape that can be used by posterior.log_prob(). Checks shapes of $\theta$ and $x$ and then repeats $x$ as often as there were batch elements in $\theta$. Args: theta: Parameters $\theta$. x: Conditioning context for posterior $p(\theta|x)$. If not provided, fall back onto an `x_o` if previously provided for multi-round training, or to another default if set later for convenience, see `.set_default_x()`. Returns: ($\theta$, $x$) with the same batch dimension, where $x$ is repeated as often as there were batch elements in $\theta$ originally. """ theta = ensure_theta_batched(torch.as_tensor(theta)) # Select and check x to condition on. x = atleast_2d_float32_tensor(self._x_else_default_x(x)) self._ensure_single_x(x) self._ensure_x_consistent_with_default_x(x) # Repeat `x` in case of evaluation on multiple `theta`. This is needed below in # when calling nflows in order to have matching shapes of theta and context x # at neural network evaluation time. x = self._match_x_with_theta_batch_shape(x, theta) return theta, x
def np_potential(self, theta: np.ndarray) -> ScalarFloat: r"""Return posterior theta log prob. $p(\theta|x)$, $-\infty$ if outside prior." Args: theta: Parameters $\theta$, batch dimension 1. Returns: Posterior log probability $\log(p(\theta|x))$. """ theta = torch.as_tensor(theta, dtype=torch.float32) theta = ensure_theta_batched(theta) num_batch = theta.shape[0] x_batched = ensure_x_batched(self.x) # Repeat x over batch dim to match theta batch, accounting for multi-D x. x_repeated = x_batched.repeat(num_batch, *(1 for _ in range(x_batched.ndim - 1))) with torch.set_grad_enabled(False): target_log_prob = self.posterior_nn.log_prob( inputs=theta.to(self.x.device), context=x_repeated, ) is_within_prior = torch.isfinite(self.prior.log_prob(theta)) target_log_prob[~is_within_prior] = -float("Inf") return target_log_prob
def log_prob( self, theta: Tensor, x: Optional[Tensor] = None, track_gradients: bool = False, ) -> Tensor: r"""Returns the log-probability of theta under the variational posterior. Args: theta: Parameters track_gradients: Whether the returned tensor supports tracking gradients. This can be helpful for e.g. sensitivity analysis but increases memory consumption. Returns: `len($\theta$)`-shaped log-probability. """ x = self._x_else_default_x(x) if self._trained_on is None or (x != self._trained_on).all(): raise AttributeError( f"The variational posterior was not fit using observation {x}.\ Please train.") with torch.set_grad_enabled(track_gradients): theta = ensure_theta_batched(torch.as_tensor(theta)) return self.q.log_prob(theta)
def __call__(self, theta: Tensor, track_gradients: bool = True) -> Tensor: r"""Returns the potential for posterior-based methods. Args: theta: The parameter set at which to evaluate the potential function. track_gradients: Whether to track the gradients. Returns: The potential. """ theta = ensure_theta_batched(torch.as_tensor(theta)) theta, x_repeated = match_theta_and_x_batch_shapes(theta, self.x_o) theta, x_repeated = theta.to(self.device), x_repeated.to(self.device) with torch.set_grad_enabled(track_gradients): posterior_log_prob = self.posterior_estimator.log_prob( theta, context=x_repeated) # Force probability to be zero outside prior support. in_prior_support = within_support(self.prior, theta) posterior_log_prob = torch.where( in_prior_support, posterior_log_prob, torch.tensor(float("-inf"), dtype=torch.float32, device=self.device), ) return posterior_log_prob
def pyro_potential( self, theta: Dict[str, Tensor], track_gradients: bool = False ) -> Tensor: r"""Return potential for Pyro sampler. Note: for Pyro this is the negative unnormalized posterior log prob. Args: theta: Parameters $\theta$. The tensor's shape will be (1, shape_of_single_theta) if running a single chain or just (shape_of_single_theta) for multiple chains. Returns: Potential $-(\log r(x_o, \theta) + \log p(\theta))$. """ theta = next(iter(theta.values())) # Theta and x should have shape (1, dim). theta = ensure_theta_batched(theta) log_ratio = RatioBasedPosterior._log_ratios_over_trials( self.x, theta.to(self.device), self.classifier, track_gradients=track_gradients, ) return -(log_ratio.cpu() + self.prior.log_prob(theta))
def posterior_potential( self, theta: np.array, track_gradients: bool = False ) -> ScalarFloat: """Returns the unnormalized posterior log-probability. This is the potential used in the numpy slice sampler and in rejection sampling. Args: theta: Parameters $\theta$, batch dimension 1. Returns: Posterior log probability of theta. """ theta = torch.as_tensor(theta, dtype=torch.float32) theta = ensure_theta_batched(theta) log_ratio = RatioBasedPosterior._log_ratios_over_trials( self.x, theta.to(self.device), self.classifier, track_gradients=track_gradients, ) # Notice opposite sign to pyro potential. return log_ratio.cpu() + self.prior.log_prob(theta)
def transformed_potential( theta: Union[Tensor, np.ndarray], potential_fn: Callable, theta_transform: torch_tf.Transform, device: str, track_gradients: bool = False, ) -> Tensor: """Return potential after a transformation by adding the log-abs-determinant. In addition, this function takes care of moving the parameters to the correct device. Args: theta: Parameters $\theta$ in transformed space. potential_fn: Potential function. theta_transform: Transformation applied before evaluating the `potential_fn` device: The device to which to move the parameters before evaluation. track_gradients: Whether or not to track the gradients of the `potential_fn` evaluation. """ # Device is the same for net and prior. transformed_theta = ensure_theta_batched( torch.as_tensor(theta, dtype=torch.float32)).to(device) # Transform `theta` from transformed (i.e. unconstrained) to untransformed # space. theta = theta_transform.inv(transformed_theta) # type: ignore log_abs_det = theta_transform.log_abs_det_jacobian(theta, transformed_theta) posterior_potential = potential_fn(theta, track_gradients=track_gradients) posterior_potential_transformed = posterior_potential - log_abs_det return posterior_potential_transformed
def log_prob( self, theta: Tensor, x: Optional[Tensor] = None, norm_posterior_snpe: bool = True, track_gradients: bool = False, ) -> Tensor: r"""Return posterior log probability $\log p(\theta|x)$. Args: theta: Parameters $\theta$. x: Conditioning context for posterior $p(\theta|x)$. If not provided, fall back onto an `x_o` if previously provided for multi-round training, or to another default if set later for convenience, see `.set_default_x()`. norm_posterior_snpe: Whether to enforce a normalized posterior density when using SNPE. Renormalization of the posterior is useful when some probability falls out or leaks out of the prescribed prior support. The normalizing factor is calculated via rejection sampling, so if you need speedier but unnormalized log posterior estimates set here `norm_posterior_snpe=False`. The returned log posterior is set to -∞ outside of the prior support regardless of this setting. track_gradients: Whether the returned tensor supports tracking gradients. This can be helpful for e.g. sensitivity analysis, but increases memory consumption. Returns: `(len(θ),)`-shaped log posterior probability $\log p(\theta|x)$ for θ in the support of the prior, -∞ (corresponding to 0 probability) outside. """ # TODO Train exited here, entered after sampling? self.net.eval() theta = ensure_theta_batched(torch.as_tensor(theta)) # Select and check x to condition on. x = atleast_2d_float32_tensor(self._x_else_default_x(x)) self._ensure_single_x(x) self._ensure_x_consistent_with_default_x(x) self._warn_if_posterior_was_focused_on_different_x(x) # Repeat `x` in case of evaluation on multiple `theta`. This is needed below in # when calling nflows in order to have matching shapes of theta and context x # at neural network evaluation time. x = self._match_x_with_theta_batch_shape(x, theta) try: log_prob_fn = getattr(self, f"_log_prob_{self._method_family}") except AttributeError: raise ValueError( f"{self._method_family} cannot evaluate probabilities.") with torch.set_grad_enabled(track_gradients): if self._method_family == "snpe": return log_prob_fn(theta, x, norm_posterior=norm_posterior_snpe) else: return log_prob_fn(theta, x)
def __init__( self, transform: torch_tf.Transform, condition: Tensor, dims_to_sample: List[int], ) -> None: super().__init__() # type: ignore self.transform = transform self.condition = ensure_theta_batched(condition) self.dims_to_sample = dims_to_sample
def _compute_covariance(probs: Tensor, limits: Tensor, f: Callable = lambda x, y: x * y) -> Tensor: """ Return the covariance between two RVs from evaluations of their pdf on a grid. The function computes the covariance as: Cov(X,Y) = E[X*Y] - E[X] * E[Y] In the more general case, when using a different function `f`, it returns: E[f(X,Y)] - f(E[X], E[Y]) By using different function `f`, this function can be also deal with more than two dimensions, but this has not been tested. Lastly, this function can also compute the variance of a 1D distribution. In that case, `probs` will be a vector, and f would be: f = lambda x: x**2: Var(X,Y) = E[X**2] - E[X]**2 Args: probs: Matrix of evaluations of a 2D density. limits: Limits within which the entries of the matrix are evenly spaced. f: The operation to be applied to the expected values, usually just the product. Returns: Covariance. """ probs = ensure_theta_batched(probs) limits = ensure_theta_batched(limits) # Compute E[X*Y]. expected_value_of_joint = _expected_value_f_of_x(probs, limits, f) # Compute E[X] * E[Y]. expected_values_of_marginals = [ _expected_value_f_of_x(prob.unsqueeze(0), lim.unsqueeze(0)) for prob, lim in zip(_calc_marginals(probs, limits), limits) ] return expected_value_of_joint - f(*expected_values_of_marginals)
def log_likelihood(self, theta: Tensor, track_gradients: bool = False) -> Tensor: """Return log likelihood of fixed data given a batch of parameters.""" log_likelihoods = LikelihoodBasedPosterior._log_likelihoods_over_trials( x=self.x, theta=ensure_theta_batched(theta).to(self.device), net=self.likelihood_nn, track_gradients=track_gradients, ) return log_likelihoods
def log_prob( self, theta: Tensor, norm_restricted_prior: bool = True, track_gradients: bool = False, prior_acceptance_params: Optional[dict] = None, ) -> Tensor: r"""Returns the log-probability of the restricted prior. Args: theta: Parameters $\theta$. norm_restricted_prior: Whether to enforce a normalized restricted prior density. The normalizing factor is calculated via rejection sampling, so if you need speedier but unnormalized log probability estimates set here `norm_restricted_prior=False`. The returned log probability is set to -∞ outside of the restriceted prior support regardless of this setting. track_gradients: Whether the returned tensor supports tracking gradients. This can be helpful for e.g. sensitivity analysis, but increases memory consumption. prior_acceptance_params: A `dict` of keyword arguments to override the default values of `prior_acceptance()`. Possible options are: `num_rejection_samples`, `force_update`, `show_progress_bars`, and `rejection_sampling_batch_size`. These parameters only have an effect if `norm_restricted_prior=True`. Returns: `(len(θ),)`-shaped log probability for θ in the support of the restricted prior, -∞ (corresponding to 0 probability) outside. """ theta = ensure_theta_batched(torch.as_tensor(theta)) with torch.set_grad_enabled(track_gradients): # Evaluate on device, move back to cpu for comparison with prior. prior_log_prob = self._prior.log_prob(theta) accepted_by_classifer = self.predict(theta) masked_log_prob = torch.where( accepted_by_classifer.bool(), prior_log_prob, torch.tensor(float("-inf"), dtype=torch.float32), ) if prior_acceptance_params is None: prior_acceptance_params = dict() # use defaults log_factor = (torch.log( self.prior_acceptance(**prior_acceptance_params)) if norm_restricted_prior else 0) return masked_log_prob - log_factor
def test_ensure_batch_dim(): # test if batch dimension is added when parameter is ndim==1 t1 = torch.tensor([0.0, -1.0, 1.0]) t2 = torchutils.ensure_theta_batched(t1) assert t2.ndim == 2 # test if batch dimension is added when observation is ndim==1 t1 = torch.tensor([0.0, -1.0, 1.0]) t2 = torchutils.ensure_x_batched(t1) assert t2.ndim == 2 # then test if batch dimension is added when observation is ndim==2, e.g. an image t1 = torch.tensor([[1, 2, 3], [1, 2, 3]]) t2 = torchutils.ensure_x_batched(t1) assert t2.ndim == 3
def potential(self, theta: Tensor, x: Optional[Tensor] = None, track_gradients: bool = False) -> Tensor: r"""Evaluates $\theta$ under the potential that is used to sample the posterior. The potential is the unnormalized log-probability of $\theta$ under the posterior. Args: theta: Parameters $\theta$. track_gradients: Whether the returned tensor supports tracking gradients. This can be helpful for e.g. sensitivity analysis, but increases memory consumption. """ self.potential_fn.set_x(self._x_else_default_x(x)) theta = ensure_theta_batched(torch.as_tensor(theta)) return self.potential_fn(theta.to(self._device), track_gradients=track_gradients)
def np_potential(self, theta: np.array) -> ScalarFloat: r"""Return posterior log prob. of theta $p(\theta|x)$" Args: theta: Parameters $\theta$, batch dimension 1. Returns: Posterior log probability of the theta, $-\infty$ if impossible under prior. """ theta = torch.as_tensor(theta, dtype=torch.float32) theta = ensure_theta_batched(theta) num_batch = theta.shape[0] x = ensure_x_batched(self.x).repeat(num_batch, 1) with torch.set_grad_enabled(False): log_likelihood = self.likelihood_nn.log_prob(inputs=x, context=theta) # Notice opposite sign to pyro potential. return log_likelihood + self.prior.log_prob(theta)
def np_potential(self, theta: np.array) -> ScalarFloat: """Return potential for Numpy slice sampler." Args: theta: Parameters $\theta$, batch dimension 1. Returns: Posterior log probability of theta. """ theta = torch.as_tensor(theta, dtype=torch.float32) # Theta and x should have shape (1, dim). theta = ensure_theta_batched(theta) x = ensure_x_batched(self.x) log_ratio = self.classifier( torch.cat((theta, x), dim=1).reshape(1, -1)) # Notice opposite sign to pyro potential. return log_ratio + self.prior.log_prob(theta)
def np_potential(self, theta: np.array) -> ScalarFloat: """Return potential for Numpy slice sampler." Args: theta: Parameters $\theta$, batch dimension 1. Returns: Posterior log probability of theta. """ theta = torch.as_tensor(theta, dtype=torch.float32) theta = ensure_theta_batched(theta) num_batch = theta.shape[0] x = ensure_x_batched(self.x).repeat(num_batch, 1) with torch.set_grad_enabled(False): log_ratio = self.classifier(torch.cat((theta, x), dim=1)).reshape(-1) # Notice opposite sign to pyro potential. return log_ratio + self.prior.log_prob(theta)
def pyro_potential(self, theta: Dict[str, Tensor]) -> Tensor: r"""Return potential for Pyro sampler. Args: theta: Parameters $\theta$. The tensor's shape will be (1, shape_of_single_theta) if running a single chain or just (shape_of_single_theta) for multiple chains. Returns: Potential $-(\log r(x_o, \theta) + \log p(\theta))$. """ theta = next(iter(theta.values())) # Theta and x should have shape (1, dim). theta = ensure_theta_batched(theta) x = ensure_x_batched(self.x) log_ratio = self.classifier([theta.to(x.device), x]).cpu() return -(log_ratio + self.prior.log_prob(theta))
def __init__( self, potential_fn_provider: Callable, condition: Tensor, dims_to_sample: List[int], ): """ Args: potential_fn_provider: Creates potential function of unconditional posterior. condition: Parameter set that all dimensions not specified in `dims_to_sample` will be fixed to. Should contain dim_theta elements, i.e. it could e.g. be a sample from the posterior distribution. The entries at all `dims_to_sample` will be ignored. dims_to_sample: Which dimensions to sample from. The dimensions not specified in `dims_to_sample` will be fixed to values given in `condition`. """ self.potential_fn_provider = potential_fn_provider self.condition = ensure_theta_batched(condition) self.dims_to_sample = dims_to_sample
def np_potential(self, theta: np.ndarray) -> ScalarFloat: r"""Return posterior theta log prob. $p(\theta|x)$, $-\infty$ if outside prior." Args: theta: Parameters $\theta$, batch dimension 1. Returns: Posterior log probability $\log(p(\theta|x))$. """ theta = torch.as_tensor(theta, dtype=torch.float32) theta = ensure_theta_batched(theta) num_batch = theta.shape[0] x = ensure_x_batched(self.x).repeat(num_batch, 1) with torch.set_grad_enabled(False): target_log_prob = self.posterior_nn.log_prob( inputs=theta, context=self.x, ) is_within_prior = torch.isfinite(self.prior.log_prob(theta)) target_log_prob[~is_within_prior] = -float("Inf") return target_log_prob
def __call__(self, theta: Tensor, track_gradients: bool = True) -> Tensor: r"""Returns the potential for posterior-based methods. Args: theta: The parameter set at which to evaluate the potential function. track_gradients: Whether to track the gradients. Returns: The potential. """ theta = ensure_theta_batched(torch.as_tensor(theta)) theta = theta.to(self.device) log_probs = [ fn(theta, track_gradients=track_gradients) for fn in self.potential_fns ] log_probs = torch.vstack(log_probs) ensemble_log_probs = torch.logsumexp( torch.log(self._weights.reshape(-1, 1)).expand_as(log_probs) + log_probs, dim=0, ) return ensemble_log_probs
def conditional_corrcoeff( density: Any, limits: Tensor, condition: Tensor, subset: Optional[List[int]] = None, resolution: int = 50, warn_about_deprecation: bool = True, ) -> Tensor: r""" Returns the conditional correlation matrix of a distribution. To compute the conditional distribution, we condition all but two parameters to values from `condition`, and then compute the Pearson correlation coefficient $\rho$ between the remaining two parameters under the distribution `density`. We do so for any pair of parameters specified in `subset`, thus creating a matrix containing conditional correlations between any pair of parameters. If `condition` is a batch of conditions, this function computes the conditional correlation matrix for each one of them and returns the mean. Args: density: Probability density function with `.log_prob()` function. limits: Limits within which to evaluate the `density`. condition: Values to condition the `density` on. If a batch of conditions is passed, we compute the conditional correlation matrix for each of them and return the average conditional correlation matrix. subset: Evaluate the conditional distribution only on a subset of dimensions. If `None` this function uses all dimensions. resolution: Number of grid points on which the conditional distribution is evaluated. A higher value increases the accuracy of the estimated correlation but also increases the computational cost. warn_about_deprecation: With sbi v0.15.0, we depracated the import of this function from `sbi.utils`. Instead, it should be imported from `sbi.analysis`. Returns: Average conditional correlation matrix of shape either `(num_dim, num_dim)` or `(len(subset), len(subset))` if `subset` was specified. """ if warn_about_deprecation: warn( "Importing `conditional_corrcoeff` from `sbi.utils` is deprecated since " "sbi v0.15.0. Instead, use " "`from sbi.analysis import conditional_corrcoeff`.") condition = ensure_theta_batched(condition) if subset is None: subset = range(condition.shape[1]) correlation_matrices = [] for cond in condition: correlation_matrices.append( torch.stack([ _compute_corrcoeff( eval_conditional_density( density, cond, limits, dim1=dim1, dim2=dim2, resolution=resolution, ), limits[[dim1, dim2]], ) for dim1 in subset for dim2 in subset if dim1 < dim2 ])) average_correlations = torch.mean(torch.stack(correlation_matrices), dim=0) # `average_correlations` is still a vector containing the upper triangular entries. # Below, assemble them into a matrix: av_correlation_matrix = torch.zeros((len(subset), len(subset))) triu_indices = torch.triu_indices(row=len(subset), col=len(subset), offset=1) av_correlation_matrix[triu_indices[0], triu_indices[1]] = average_correlations # Make the matrix symmetric by copying upper diagonal to lower diagonal. av_correlation_matrix = torch.triu(av_correlation_matrix) + torch.tril( av_correlation_matrix.T) av_correlation_matrix.fill_diagonal_(1.0) return av_correlation_matrix
def log_prob( self, theta: Tensor, x: Optional[Tensor] = None, norm_posterior: bool = True, track_gradients: bool = False, leakage_correction_params: Optional[dict] = None, ) -> Tensor: r"""Returns the log-probability of the posterior $p(\theta|x)$. Args: theta: Parameters $\theta$. norm_posterior: Whether to enforce a normalized posterior density. Renormalization of the posterior is useful when some probability falls out or leaks out of the prescribed prior support. The normalizing factor is calculated via rejection sampling, so if you need speedier but unnormalized log posterior estimates set here `norm_posterior=False`. The returned log posterior is set to -∞ outside of the prior support regardless of this setting. track_gradients: Whether the returned tensor supports tracking gradients. This can be helpful for e.g. sensitivity analysis, but increases memory consumption. leakage_correction_params: A `dict` of keyword arguments to override the default values of `leakage_correction()`. Possible options are: `num_rejection_samples`, `force_update`, `show_progress_bars`, and `rejection_sampling_batch_size`. These parameters only have an effect if `norm_posterior=True`. Returns: `(len(θ),)`-shaped log posterior probability $\log p(\theta|x)$ for θ in the support of the prior, -∞ (corresponding to 0 probability) outside. """ x = self._x_else_default_x(x) # TODO Train exited here, entered after sampling? self.posterior_estimator.eval() theta = ensure_theta_batched(torch.as_tensor(theta)) theta_repeated, x_repeated = match_theta_and_x_batch_shapes(theta, x) with torch.set_grad_enabled(track_gradients): # Evaluate on device, move back to cpu for comparison with prior. unnorm_log_prob = self.posterior_estimator.log_prob( theta_repeated, context=x_repeated) # Force probability to be zero outside prior support. in_prior_support = within_support(self.prior, theta_repeated) masked_log_prob = torch.where( in_prior_support, unnorm_log_prob, torch.tensor(float("-inf"), dtype=torch.float32, device=self._device), ) if leakage_correction_params is None: leakage_correction_params = dict() # use defaults log_factor = (log( self.leakage_correction(x=x, **leakage_correction_params)) if norm_posterior else 0) return masked_log_prob - log_factor
def eval_conditional_density( density: Any, condition: Tensor, limits: Tensor, dim1: int, dim2: int, resolution: int = 50, eps_margins1: Union[Tensor, float] = 1e-32, eps_margins2: Union[Tensor, float] = 1e-32, warn_about_deprecation: bool = True, ) -> Tensor: r""" Return the unnormalized conditional along `dim1, dim2` given parameters `condition`. We compute the unnormalized conditional by evaluating the joint distribution: $p(x1 | x2) = p(x1, x2) / p(x2) \propto p(x1, x2)$ Args: density: Probability density function with `.log_prob()` method. condition: Parameter set that all dimensions other than dim1 and dim2 will be fixed to. Should be of shape (1, dim_theta), i.e. it could e.g. be a sample from the posterior distribution. The entries at `dim1` and `dim2` will be ignored. limits: Bounds within which to evaluate the density. Shape (dim_theta, 2). dim1: First dimension along which to evaluate the conditional. dim2: Second dimension along which to evaluate the conditional. resolution: Resolution of the grid along which the conditional density is evaluated. eps_margins1: We will evaluate the posterior along `dim1` at `limits[0]+eps_margins` until `limits[1]-eps_margins`. This avoids evaluations potentially exactly at the prior bounds. eps_margins2: We will evaluate the posterior along `dim2` at `limits[0]+eps_margins` until `limits[1]-eps_margins`. This avoids evaluations potentially exactly at the prior bounds. warn_about_deprecation: With sbi v0.15.0, we depracated the import of this function from `sbi.utils`. Instead, it should be imported from `sbi.analysis`. Returns: Conditional probabilities. If `dim1 == dim2`, this will have shape (resolution). If `dim1 != dim2`, it will have shape (resolution, resolution). """ if warn_about_deprecation: warn( "Importing `eval_conditional_density` from `sbi.utils` is deprecated since " "sbi v0.15.0. Instead, use " "`from sbi.analysis import eval_conditional_density`.") condition = ensure_theta_batched(condition) theta_grid_dim1 = torch.linspace( float(limits[dim1, 0] + eps_margins1), float(limits[dim1, 1] - eps_margins1), resolution, ) theta_grid_dim2 = torch.linspace( float(limits[dim2, 0] + eps_margins2), float(limits[dim2, 1] - eps_margins2), resolution, ) if dim1 == dim2: repeated_condition = condition.repeat(resolution, 1) repeated_condition[:, dim1] = theta_grid_dim1 log_probs_on_grid = density.log_prob(repeated_condition) else: repeated_condition = condition.repeat(resolution**2, 1) repeated_condition[:, dim1] = theta_grid_dim1.repeat(resolution) repeated_condition[:, dim2] = torch.repeat_interleave( theta_grid_dim2, resolution) log_probs_on_grid = density.log_prob(repeated_condition) log_probs_on_grid = torch.reshape(log_probs_on_grid, (resolution, resolution)) # Subtract maximum for numerical stability. return torch.exp(log_probs_on_grid - torch.max(log_probs_on_grid))
def eval_conditional_density( density: Any, condition: Tensor, limits: Tensor, dim1: int, dim2: int, resolution: int = 50, eps_margins1: Union[Tensor, float] = 1e-32, eps_margins2: Union[Tensor, float] = 1e-32, return_raw_log_prob: bool = False, ) -> Tensor: r"""Return the unnormalized conditional along `dim1, dim2` given `condition`. We compute the unnormalized conditional by evaluating the joint distribution: $p(x1 | x2) = p(x1, x2) / p(x2) \propto p(x1, x2)$ The joint distribution is evaluated on an evenly spaced grid defined by the `limits`. Args: density: Probability density function with `.log_prob()` method. condition: Parameter set that all dimensions other than dim1 and dim2 will be fixed to. Should be of shape (1, dim_theta), i.e. it could e.g. be a sample from the posterior distribution. The entries at `dim1` and `dim2` will be ignored. limits: Bounds within which to evaluate the density. Shape (dim_theta, 2). dim1: First dimension along which to evaluate the conditional. dim2: Second dimension along which to evaluate the conditional. resolution: Resolution of the grid along which the conditional density is evaluated. eps_margins1: We will evaluate the posterior along `dim1` at `limits[0]+eps_margins` until `limits[1]-eps_margins`. This avoids evaluations potentially exactly at the prior bounds. eps_margins2: We will evaluate the posterior along `dim2` at `limits[0]+eps_margins` until `limits[1]-eps_margins`. This avoids evaluations potentially exactly at the prior bounds. return_raw_log_prob: If `True`, return the log-probability evaluated on the grid. If `False`, return the probability, scaled down by the maximum value on the grid for numerical stability (i.e. exp(log_prob - max_log_prob)). Returns: Conditional probabilities. If `dim1 == dim2`, this will have shape (resolution). If `dim1 != dim2`, it will have shape (resolution, resolution). """ condition = ensure_theta_batched(condition) theta_grid_dim1 = torch.linspace( float(limits[dim1, 0] + eps_margins1), float(limits[dim1, 1] - eps_margins1), resolution, device=condition.device, ) theta_grid_dim2 = torch.linspace( float(limits[dim2, 0] + eps_margins2), float(limits[dim2, 1] - eps_margins2), resolution, device=condition.device, ) if dim1 == dim2: repeated_condition = condition.repeat(resolution, 1) repeated_condition[:, dim1] = theta_grid_dim1 log_probs_on_grid = density.log_prob(repeated_condition) else: repeated_condition = condition.repeat(resolution**2, 1) repeated_condition[:, dim1] = theta_grid_dim1.repeat(resolution) repeated_condition[:, dim2] = torch.repeat_interleave( theta_grid_dim2, resolution) log_probs_on_grid = density.log_prob(repeated_condition) log_probs_on_grid = torch.reshape(log_probs_on_grid, (resolution, resolution)) if return_raw_log_prob: return log_probs_on_grid else: # Subtract maximum for numerical stability return torch.exp(log_probs_on_grid - torch.max(log_probs_on_grid))