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