import keras.ops as ops
import numpy as np
import math
from bayesflow.utils.serialization import serializable, serialize
from bayesflow.types import Tensor
from .constrain import Constrain
from .elementwise_transform import ElementwiseTransform
[docs]
@serializable("bayesflow.adapters")
class AsCovarianceMatrix(ElementwiseTransform):
"""
Constrains neural network predictions of a variable to a valid (symmetric,
positive definite) K x K covariance (or precision) matrix.
The unconstrained representation is a flat vector of K * (K + 1) / 2 entries
of the lower Cholesky factor of the covariance matrix.
First, the K diagonal entries, passed through a lower-bounded :py:class:`~transforms.Constrain`
transform to keep them positive, followed by the K * (K - 1) / 2 entries from the lower triangular.
Parameters
----------
cholesky : bool, optional
Whether the *constrained* side of the transform is the Cholesky factor,
rather than the full covariance matrix.
Default is False (use the full covariance matrix).
diag_kwargs : dict, optional
Keyword arguments forwarded to the :py:class:`~transforms.Constrain` transform
that constrains the diagonal to be positive.
The `lower` bound is always fixed to 0.0. All other arguments use the default
of :py:class:`~transforms.Constrain`.
Examples
--------
>>> adapter = bf.Adapter().as_covariance_matrix("cov_x", diag_kwargs={"method": "exp"})
"""
def __init__(self, *, cholesky: bool = False, diag_kwargs: dict = None):
super().__init__()
self.cholesky = cholesky
self.diag_kwargs = diag_kwargs or {}
self.diag_kwargs["lower"] = 0.0
self.diag_transform = Constrain(**self.diag_kwargs)
[docs]
def forward(self, data: Tensor, **kwargs) -> Tensor:
K = ops.shape(data)[-1]
batch_shape = tuple(ops.shape(data)[:-2])
if not self.cholesky:
data = ops.cholesky(data)
diag = ops.diagonal(data, axis1=-2, axis2=-1)
diag = self.diag_transform(diag)
# extract flattened off-diagonal
rows, cols = np.tril_indices(K, k=-1)
idx = K * rows + cols
data_flat = ops.reshape(data, batch_shape + (K * K,))
lower = ops.take(data_flat, idx, axis=-1)
return ops.concatenate([diag, lower], axis=-1)
[docs]
def inverse(self, data: Tensor, **kwargs) -> Tensor:
P = ops.shape(data)[-1]
K = (math.isqrt(1 + 8 * P) - 1) // 2
batch_shape = tuple(ops.shape(data)[:-1])
diag = self.diag_transform(data[..., :K], inverse=True)
lower = data[..., K:]
# Index map into [0, diag..., lower...] for each flattened (i, j) position
rows, cols = np.tril_indices(K, k=-1)
idx_map = np.zeros(K * K, dtype="int32") # upper triangle -> slot 0 (zero)
idx_map[np.arange(K) * (K + 1)] = 1 + np.arange(K) # diagonal
idx_map[rows * K + cols] = 1 + K + np.arange(len(rows)) # strictly lower
padded = ops.concatenate([ops.zeros_like(diag[..., :1]), diag, lower], axis=-1)
L = ops.reshape(ops.take(padded, idx_map, axis=-1), batch_shape + (K, K))
if not self.cholesky:
return ops.matmul(L, ops.swapaxes(L, -1, -2)) # L @ L^T
return L
[docs]
def get_config(self) -> dict:
config = {
"cholesky": self.cholesky,
"diag_kwargs": self.diag_kwargs,
}
return serialize(config)
[docs]
def log_det_jac(self, data: Tensor, inverse: bool = False, **kwargs) -> Tensor:
K = ops.shape(data)[-1]
if not self.cholesky:
data = ops.cholesky(data)
diag = ops.diagonal(data, axis1=-2, axis2=-1)
ldj = ops.sum(self.diag_transform.ldj(diag), axis=-1)
if not self.cholesky:
weight = ops.cast(K, diag.dtype) - ops.arange(K, dtype=diag.dtype)
chol_to_matrix_ldj = K * ops.log(ops.cast(2.0, diag.dtype)) + ops.sum(weight * ops.log(diag), axis=-1)
ldj = ldj - chol_to_matrix_ldj
if inverse:
ldj = -ldj
return self._sum_except_batch(ldj)
