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))
