# # Copyright (c) 2021 The Markovflow Contributors. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # """Module containing a multivariate Gaussian likelihood.""" from typing import Tuple import tensorflow as tf import tensorflow_probability as tfp from gpflow import default_float from gpflow.base import Parameter, TensorType from gpflow.utilities import triangular from markovflow.likelihoods.likelihoods import Likelihood, check_input_shapes from markovflow.utils import tf_scope_class_decorator [docs]@tf_scope_class_decorator class MultivariateGaussian(Likelihood): """ Represents a multivariate Gaussian likelihood. For example: .. math:: p(yᵢ | fᵢ) = 𝓝(yᵢ; fᵢ, Σ= LLᵀ) See also the documentation for the base :class:`~markovflow.likelihoods.likelihoods.Likelihood` class. """ def __init__(self, chol_covariance: TensorType): """ :param chol_covariance: A :data:`~markovflow.base.TensorType` containing the Cholesky factor of the covariance of the Gaussian noise, with shape ``[obs_dim, obs_dim]``. """ super().__init__(self.__class__.__name__) self.chol_covariance = Parameter( chol_covariance, transform=triangular(), name="chol_covariance" ) self._obs_dim = self.chol_covariance.shape[-1] tf.debugging.assert_shapes( [(self.chol_covariance, [..., self.obs_dim, self.obs_dim]),] ) @property [docs] def obs_dim(self) -> int: """ Return the dimensionality of each observation. """ return self._obs_dim [docs] def log_probability_density(self, fs: tf.Tensor, observations: tf.Tensor) -> tf.Tensor: """ Compute the log probability density :math:`log p(Y|F)`. For a multivariate Gaussian, this is :math:`log 𝓝(yᵢ; fᵢ, Σ)`. :param fs: A tensor representing a conditioning variable, with shape ``batch_shape + [num_data, obs_dim]``. :param observations: A tensor representing a conditioned variable, with shape ``batch_shape + [num_data, obs_dim]``. :return: A tensor representing :math:`log p(yᵢ | fᵢ)`, with shape ``batch_shape + [num_data]``. """ return tfp.distributions.MultivariateNormalTriL(fs, self.chol_covariance).log_prob( observations ) [docs] def variational_expectations( self, f_means: tf.Tensor, f_covariances: tf.Tensor, observations: tf.Tensor ) -> tf.Tensor: """ Calculate a variational expectation for each observation: .. math:: ∫ q(fᵢ) log p(yᵢ|fᵢ) dfᵢ ...where: * :math:`q(fᵢ) ~ N(μᵢ, Σᵢ)` * :math:`p(y |f)` is a general likelihood function For a multivariate Gaussian this is: .. math:: ∫ 𝓝(fᵢ; μᵢ, Sᵢ) log𝓝(yᵢ; fᵢ, Σ= LLᵀ) dfᵢ = -½ Tr(Σ⁻¹Sᵢ) + log𝓝(yᵢ; μᵢ, Σ) :param f_means: The marginal :math:`f` means for each state of the :class:`~markovflow.state_space_model.StateSpaceModel`, with shape ``batch_shape + [num_data, obs_dim]``. :param f_covariances: The marginal :math:`f` covariances for each state of the :class:`~markovflow.state_space_model.StateSpaceModel`, with shape ``batch_shape + [num_data, obs_dim, obs_dim]``. :param observations: The :math:`y` values at which to evaluate the log probability, with shape ``batch_shape + [num_data, obs_dim]``. :return: A tensor with shape ``batch_shape + [num_data]``. """ Y = observations check_input_shapes(f_means, Y, self.obs_dim, f_covariances=f_covariances) inv_covariance = tf.linalg.cholesky_solve( self.chol_covariance, tf.eye(self.obs_dim, dtype=default_float()) ) return -0.5 * tf.reduce_sum( inv_covariance * f_covariances, axis=[-1, -2] ) + tfp.distributions.MultivariateNormalTriL(f_means, self.chol_covariance).log_prob(Y) [docs] def predict_density( self, f_means: tf.Tensor, f_covariances: tf.Tensor, observations: tf.Tensor ) -> tf.Tensor: """ Predict a density. This calculates: .. math:: ∫ q(F) p(Y|F) dF ...of a Gaussian approximation :math:`q(F) ~ N(μ, Σ)` to the posterior density :math:`p(F|Y)`. For a multivariate Gaussian this is: .. math:: log ∫ 𝓝(fᵢ; μᵢ, Sᵢ) 𝓝(yᵢ; fᵢ, Σ= LLᵀ) dfᵢ = log 𝓝(yᵢ; μᵢ, Σ + Sᵢ) :param f_means: The marginal :math:`f` means for each state of the :class:`~markovflow.state_space_model.StateSpaceModel`, with shape ``batch_shape + [num_data, obs_dim]``. :param f_covariances: The marginal :math:`f` covariances for each state of the :class:`~markovflow.state_space_model.StateSpaceModel`, with shape ``batch_shape + [num_data, obs_dim, obs_dim]``. :param observations: The :math:`y` values at which to evaluate the log probability, with shape ``batch_shape + [num_data, obs_dim]``. """ Y = observations check_input_shapes(f_means, Y, self.obs_dim, f_covariances=f_covariances) covariance = tf.matmul(self.chol_covariance, self.chol_covariance, transpose_b=True) f_chol_covariances = tf.linalg.cholesky(f_covariances + covariance) return tfp.distributions.MultivariateNormalTriL(f_means, f_chol_covariances).log_prob(Y) [docs] def predict_mean_and_var( self, f_means: tf.Tensor, f_covariances: tf.Tensor ) -> Tuple[tf.Tensor, tf.Tensor]: """ Predict the observation means and covariances given the f-space means and covariances. That is, calculate: .. math:: p(y* | x*, x, y) = ∫ p(y* | f*) p(f* | x*, x, y) df* ...where: * `f_means` and `f_covariances` is our representation of :math:`p(f* | x*, x, y)` * :math:`p(y* | f*)` is defined by the likelihood For a multivariate Gaussian this is: .. math:: p(y* | x*, x, y) = ∫ 𝓝(f*; μ*, S*) 𝓝(y*; f*, Σ= LLᵀ) dfᵢ = 𝓝(y*; μ*, Σ + S*) :param f_means: The marginal :math:`f` means for some arbitrary predicted time points, with shape ``batch_shape + [num_data, obs_dim]``. :param f_covariances: The marginal :math:`f` covariances for some arbitrary predicted time points, with shape ``batch_shape + [num_data, obs_dim, obs_dim]``. :return: A tuple of tensors containing observation means and covariances, with respective shapes ``batch_shape + [num_time_points, obs_dim]``, ``batch_shape + [num_time_points, obs_dim, obs_dim]``. """ covariance = tf.matmul(self.chol_covariance, self.chol_covariance, transpose_b=True) return f_means, covariance + f_covariances