Source code for gmmvi.experiments.target_distributions.gmm

import math

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
import matplotlib.pyplot as plt

from gmmvi.experiments.target_distributions.logistic_regression import LNPDF
from gmmvi.models.gmm_wrapper import GmmWrapper


[docs]class GMM_LNPDF(LNPDF): """Implements a target distribution that is given by a Gaussian mixture model. Parameters ---------- target_weights: tf.Tensor of tf.float32 a one-dimensional vector of size number_of_components containing the mixture weights. target_means: tf.Tensor of tf.float32 a two-dimensional vector of size number_of_components x dimensions containing the mixture means. target_covs: tf.Tensor of tf.float32 a three-dimensional vector of size number_of_components x dimensions x dimensions containing the covariance matrices. """ def __init__(self, target_weights: tf.Tensor, target_means: tf.Tensor, target_covs: tf.Tensor): super(GMM_LNPDF, self).__init__(use_log_density_and_grad=False, safe_for_tf_graph=True) self.target_weights = tf.cast(target_weights, dtype=tf.float32) self.target_means = tf.cast(target_means, dtype=tf.float32) self.target_covs = tf.cast(target_covs, dtype=tf.float32) self.gmm = tfp.distributions.MixtureSameFamily( mixture_distribution=tfp.distributions.Categorical(logits=tf.math.log(self.target_weights)), components_distribution=tfp.distributions.MultivariateNormalTriL( loc=self.target_means, scale_tril=tf.linalg.cholesky(self.target_covs)))
[docs] def log_density(self, x): x = tf.cast(x, dtype=tf.float32) return self.gmm.log_prob(x)
[docs] def marginal_log_density(self, x, dim): """ Computes the marginal distribution along the given dimensions. Parameters ---------- x: tf.Tensor of tf.float32 a one-dimensional vector of size number_of_samples containing the samples we want to evaluate dim: an int Specifies the dimension used for constructing the marginal GMM. Returns: tf.Tensor - a one-dimensional Tensor of shape number_of_samples containing the marginal log densities. """ mixture = tfp.distributions.MixtureSameFamily( mixture_distribution=tfp.distributions.Categorical(logits=tf.math.log(self.target_weights)), components_distribution=tfp.distributions.Normal(loc=self.target_means[:, dim], scale=tf.math.sqrt(self.target_covs[:, dim, dim]))) x = tf.cast(x, dtype=tf.float32) return mixture.log_prob(x)
[docs] def get_num_dimensions(self): return len(self.target_means[0])
[docs] def can_sample(self): """ Returns: bool: We can sample from a GMM, so this method will return True. """ return True
[docs] def sample(self, n): """ Draws n samples from this GMM. Parameters: n: int The number of samples we want to draw. Returns: tf.Tensor: The sample, a tensor of size n x dimensions. """ return self.gmm.sample(n)
[docs] def expensive_metrics(self, model: GmmWrapper, samples: tf.Tensor) -> dict: """ This method computed the number of detected modes (by testing how many modes of this target distribution are close to a component in the learned model) and a figure that shows plots comparing the marginal distributions of the model with the true marginals of this target distribution. Parameters: model: :py:class:`GmmWrapper<gmmvi.models.gmm_wrapper.GmmWrapper>` The learned model that we want to evaluate for this target distribution. samples: tf.Tensor Samples that have been drawn from the model, which can be used for evaluations. Returns: dict: a dictionary containing two items (the number of detected modes, and a figure showing the plots of marginals. """ expensive_metrics = dict() x_vals = tf.linspace(-70, 70, 1000) x_vals = tf.reshape(tf.repeat(x_vals, [self.get_num_dimensions()]), [-1, self.get_num_dimensions()]) x_vals = tf.cast(x_vals, dtype=tf.float32) fig, axs = plt.subplots(4, 5) for dim in range(20): true_densities = tf.exp(self.marginal_log_density(x_vals, dim)) start = tf.reduce_min(tf.where(true_densities > 1e-4)) end = tf.reduce_max(tf.where(true_densities > 1e-4)) axs[dim // 5, dim % 5].plot(x_vals[start:end, dim], true_densities[start:end], color='b') axs[dim // 5, dim % 5].plot(x_vals[start:end, dim], tf.exp(model.marginal_log_density(x_vals[start:end], dim)), color='r') dists_to_target_means = tf.reduce_min( tf.linalg.norm(tf.expand_dims(self.target_means, 1) - tf.expand_dims(model.means, 0), axis=2), 1) detection_threshold = tf.linalg.norm(6. * tf.ones(model.num_dimensions)) num_detected = tf.where(dists_to_target_means < detection_threshold).shape[0] print(f"Found {num_detected} components.") expensive_metrics.update({"num_detected_modes": num_detected, "marginals": fig}) return expensive_metrics
def make_target(num_dimensions): """ Create a :py:class:`GMM target distribution<gmmvi.experiments.target_distributions.gmm.GMM_LNPDF>` using the same procedure as :cite:t:`Arenz2020` for initializing the weights, means and covariance matrices. Parameters: num_dimensions: int The number of dimensions of the target GMM Returns: :py:class:`GMM_LNPDF<gmmvi.experiments.target_distributions.gmm.GMM_LNPDF>`: the instantiated object """ num_true_components = 10 weights = np.ones(num_true_components) / num_true_components means = np.empty((num_true_components, num_dimensions)) covs = np.empty((num_true_components, num_dimensions, num_dimensions)) for i in range(0, num_true_components): means[i] = 100 * (np.random.random(num_dimensions) - 0.5) covs[i] = 0.1 * np.random.normal(0, num_dimensions, (num_dimensions * num_dimensions)).reshape( (num_dimensions, num_dimensions)) covs[i] = covs[i].transpose().dot(covs[i]) covs[i] += 1 * np.eye(num_dimensions) return GMM_LNPDF(weights, means, covs) def make_target_with_scale(num_dimensions, num_components, scale): num_true_components = num_components weights = np.ones(num_true_components) / num_true_components means = np.empty((num_true_components, num_dimensions)) covs = np.empty((num_true_components, num_dimensions, num_dimensions)) for i in range(0, num_true_components): means[i] = 100 * (np.random.random(num_dimensions) - 0.5) covs[i] = np.random.normal(0, np.sqrt(scale), (num_dimensions * num_dimensions)).reshape( (num_dimensions, num_dimensions)) covs[i] = covs[i].transpose().dot(covs[i]) covs[i] += 1 * np.eye(num_dimensions) # print("target_means:", means) # print("target_covs:", covs) return GMM_LNPDF(weights, means, covs) def U(theta): return np.array( [ [math.cos(theta), math.sin(theta)], [-math.sin(theta), math.cos(theta)], ] ) def make_simple_target(): pi = math.pi # weights w_true = np.array([0.5, 0.3, 0.2]) # means mu_true = np.array( [ [-2.0, -2.0], [2.0, -2.0], [0.0, 2.0], ] ) # covs cov1 = np.array([[0.5, 0.0], [0.0, 1.0]]) cov1 = U(pi / 4) @ cov1 @ np.transpose(U(pi / 4)) cov2 = np.array([[0.5, 0.0], [0.0, 1.0]]) cov2 = U(-pi / 4) @ cov2 @ np.transpose(U(-pi / 4)) cov3 = np.array([[1.0, 0.0], [0.0, 2.0]]) cov3 = U(pi / 2) @ cov3 @ np.transpose(U(pi / 2)) cov_true = np.stack([cov1, cov2, cov3], axis=0) # generate target dist target_dist = GMM_LNPDF( target_weights=w_true, target_means=mu_true, target_covs=cov_true, ) return target_dist def make_star_target(num_components): # Source: Lin et al. K = num_components ## weights w_true = np.ones((K,)) / K ## means and precs # first component mus = [np.array([1.5, 0.0])] precs = [np.diag([1.0, 100.0])] # other components are generated through rotation theta = 2 * math.pi / K for _ in range(K - 1): mus.append(U(theta) @ mus[-1]) precs.append(U(theta) @ precs[-1] @ np.transpose(U(theta))) assert len(w_true) == len(mus) == len(precs) == K mu_true = np.stack(mus, axis=0) prec_true = np.stack(precs, axis=0) cov_true = np.linalg.inv(prec_true) # generate target dist target_dist = GMM_LNPDF( target_weights=w_true, target_means=mu_true, target_covs=cov_true, ) return target_dist