import math
import keras
from bayesflow.types import Shape, Tensor
from bayesflow.links import CholeskyFactor
from bayesflow.utils.keras_utils import resolve_seed
from bayesflow.utils.serialization import serializable
from .parametric_distribution_score import ParametricDistributionScore
[docs]
@serializable("bayesflow.scoring_rules", disable_module_check=True)
class MvNormalScore(ParametricDistributionScore):
r""":math:`S(\hat p_{\mu, \Sigma}, \theta; k) = -\log( \mathcal N (\theta; \mu, \Sigma))`
Scores a predicted mean and lower-triangular Cholesky factor :math:`L` of the precision matrix :math:`P`
with the log-score of the probability of the materialized value. The precision matrix is
the inverse of the covariance matrix, :math:`L^T L = P = \Sigma^{-1}`.
"""
NOT_TRANSFORMING_LIKE_VECTOR_WARNING = ("precision_cholesky_factor",)
"""
Marks head for precision matrix Cholesky factor as an exception for adapter transformations.
This variable contains names of prediction heads that should lead to a warning when the adapter is applied
in inverse direction to them.
For more information see :py:class:`ScoringRule`.
"""
TRANSFORMATION_TYPE: dict[str, str] = {"precision_cholesky_factor": "right_side_scale_inverse"}
"""
Marks precision Cholesky factor head to handle de-standardization appropriately.
See :py:class:`bayesflow.networks.Standardization` for more information on supported de-standardization options.
For the mean head the default ("location_scale") is not overridden.
"""
def __init__(self, dim: int = None, links: dict = None, **kwargs):
super().__init__(links=links, **kwargs)
self.dim = dim
self.links = links or {"precision_cholesky_factor": CholeskyFactor()}
self.config = {"dim": dim}
[docs]
def get_config(self):
base_config = super().get_config()
return base_config | self.config
[docs]
def get_head_shapes_from_target_shape(self, target_shape: Shape) -> dict[str, Shape]:
self.dim = target_shape[-1]
return dict(mean=(self.dim,), precision_cholesky_factor=(self.dim, self.dim))
[docs]
def log_prob(self, x: Tensor, mean: Tensor, precision_cholesky_factor: Tensor) -> Tensor:
"""
Compute the log probability density of a multivariate Gaussian distribution.
This function calculates the log probability density for each sample in `x` under a
multivariate Gaussian distribution with the given `mean` and `precision_cholesky_factor`.
The computation includes the determinant of the precision matrix, its inverse, and the quadratic
form in the exponential term of the Gaussian density function.
Parameters
----------
x : Tensor
A tensor of input samples for which the log probability density is computed.
The shape should be compatible with broadcasting against `mean`.
mean : Tensor
A tensor representing the mean of the multivariate Gaussian distribution.
precision_cholesky_factor : Tensor
A tensor representing the lower-triangular Cholesky factor of the precision matrix
of the multivariate Gaussian distribution.
Returns
-------
Tensor
A tensor containing the log probability densities for each sample in `x` under the
given Gaussian distribution.
"""
diff = x - mean
# Compute log determinant, exploiting Cholesky factors
log_det_covariance = -2 * keras.ops.sum(
keras.ops.log(keras.ops.diagonal(precision_cholesky_factor, axis1=1, axis2=2)), axis=1
)
# Compute the quadratic term in the exponential of the multivariate Gaussian from Cholesky factors
# diff^T * precision_cholesky_factor^T * precision_cholesky_factor * diff
quadratic_term = keras.ops.einsum(
"...i,...ji,...jk,...k->...", diff, precision_cholesky_factor, precision_cholesky_factor, diff
)
# Compute the log probability density
log_prob = -0.5 * (self.dim * keras.ops.log(2 * math.pi) + log_det_covariance + quadratic_term)
return log_prob
[docs]
def sample(
self,
batch_shape: Shape,
mean: Tensor,
precision_cholesky_factor: Tensor,
seed: int | keras.random.SeedGenerator | None = None,
) -> Tensor:
"""
Generate samples from a multivariate Gaussian distribution.
Draws one sample per parameter set by transforming standard normal samples
using the Cholesky factor of the covariance matrix.
Parameters
----------
batch_shape : Shape
Ignored. Retained for interface compatibility. The output shape
is determined by the shape of the parameters.
mean : Tensor
A tensor representing the mean of the multivariate Gaussian distribution.
Shape: ``(..., D)``.
precision_cholesky_factor : Tensor
A tensor representing the lower-triangular Cholesky factor of the precision matrix
of the multivariate Gaussian distribution.
Shape: ``(..., D, D)``.
seed : int, keras.random.SeedGenerator, or None, optional
Seed for reproducible sampling.
Returns
-------
Tensor
Samples with the same shape as ``mean``.
"""
covariance_cholesky_factor = keras.ops.inv(precision_cholesky_factor)
normal_samples = keras.random.normal(keras.ops.shape(mean), seed=resolve_seed(seed))
scaled_normal = keras.ops.einsum("...ij,...j->...i", covariance_cholesky_factor, normal_samples)
samples = mean + scaled_normal
return samples
