Source code for gmmvi.models.full_cov_gmm

import tensorflow as tf
from gmmvi.models.gmm import GMM
from math import pi


[docs]class FullCovGMM(GMM): """ A Gaussian mixture model with full covariance matrices. Parameters: weights: tf.Tensor a one-dimensional tensor containing the initial weights of the GMM. means: tf.Tensor a two-dimensional tensor containing the component means. covs: tf.Tensor a three-dimensional tensor containing the component covariance matrices. """ def __init__(self, weights: tf.Tensor, means: tf.Tensor, covs: tf.Tensor): self.diagonal_covs = False num_dimensions = len(means[0]) log_weights = tf.Variable(tf.math.log(weights), shape=[None], dtype=tf.float32) chol_covs = tf.Variable(tf.stack([tf.linalg.cholesky(cov) for cov in covs]), shape=[None, num_dimensions, num_dimensions], dtype=tf.float32) means = tf.Variable(tf.convert_to_tensor(means), shape=[None, num_dimensions], dtype=tf.float32) super(FullCovGMM, self).__init__(log_weights, means, chol_covs) @property def covs(self) -> tf.Tensor: return self.chol_cov @ tf.transpose(self.chol_cov, [0, 2, 1])
[docs] def gaussian_entropy(self, chol: tf.Tensor) -> tf.Tensor: return 0.5 * self.num_dimensions * (tf.math.log(2 * pi) + 1) + tf.reduce_sum(tf.math.log(tf.linalg.tensor_diag_part(chol)))
[docs] def sample_from_component(self, index: int, num_samples: int) -> tf.Tensor: return tf.transpose(tf.expand_dims(self.means[index], axis=-1) + self.chol_cov[index] @ tf.random.normal((self.num_dimensions, num_samples), mean=0., stddev=1.))
[docs] def component_log_density(self, index: int, samples: tf.Tensor) -> tf.Tensor: diffs = samples - self.means[index] sqrts = tf.linalg.triangular_solve(self.chol_cov[index], tf.transpose(diffs)) mahalas = - 0.5 * tf.reduce_sum(sqrts * sqrts, axis=0) const_parts = - 0.5 * tf.reduce_sum(tf.math.log(tf.square(tf.linalg.diag_part(self.chol_cov[index])))) \ - 0.5 * self.num_dimensions * tf.math.log(2 * pi) return mahalas + const_parts
[docs] def component_marginal_log_densities(self, samples: tf.Tensor, dim: int) -> tf.Tensor: diffs = tf.expand_dims(samples[:, dim], 0) - tf.expand_dims(self.means[:, dim], 1) mahalas = - 0.5 * diffs * diffs / tf.expand_dims(self.covs[:, dim, dim], 1) const_parts = - 0.5 * tf.math.log(self.covs[:, dim, dim]) \ - 0.5 * tf.math.log(2 * pi) return mahalas + tf.expand_dims(const_parts, axis=1)
[docs] def component_log_densities(self, samples: tf.Tensor) -> tf.Tensor: diffs = tf.expand_dims(samples, 0) - tf.expand_dims(self.means, 1) sqrts = tf.linalg.triangular_solve(self.chol_cov, tf.transpose(diffs, [0, 2, 1])) mahalas = - 0.5 * tf.reduce_sum(sqrts * sqrts, axis=1) const_parts = - 0.5 * tf.reduce_sum(tf.math.log(tf.square(tf.linalg.diag_part(self.chol_cov))), axis=1) \ - 0.5 * self.num_dimensions * tf.math.log(2 * pi) return mahalas + tf.expand_dims(const_parts, axis=1)
[docs] def add_component(self, initial_weight: tf.Tensor, initial_mean: tf.Tensor, initial_cov: tf.Tensor): self.means.assign(tf.concat((self.means, tf.expand_dims(initial_mean, axis=0)), axis=0)) self.chol_cov.assign( tf.concat((self.chol_cov, tf.expand_dims(tf.linalg.cholesky(initial_cov), axis=0)), axis=0)) self.replace_weights(tf.concat((self.log_weights, tf.expand_dims(tf.math.log(initial_weight), axis=0)), axis=0))