Skip to content

Mathematical functions

VAE.utils.math

Collection of mathematical functions used in the VAE.

VAE.utils.math.log_density_gaussian 

log_density_gaussian(sample, mean, logvar, all_combinations=False, axis=0)

Calculates the sampled log density of a Gaussian.

Version for Tensor arrays as inputs.

Parameters:

  • sample (Tensor) –

    Sample at which to compute the density.

  • mean (Tensor) –

    Mean of Gaussian distribution.

  • logvar (Tensor) –

    Log variance of Gaussian distribution.

  • all_combinations (bool, default: False ) –

    Whether to return values for all combinations of batch pairs of sample and mean.

  • axis (int, default: 0 ) –

    The dimension over which all combinations are computed. Default is 0 meaning the batch dimension.

Returns:

  • Tensor –

    Tensor

Source code in VAE/utils/math.py 
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
def log_density_gaussian(sample: tf.Tensor,
                         mean: tf.Tensor,
                         logvar: tf.Tensor,
                         all_combinations: bool = False,
                         axis: int = 0) -> tf.Tensor:
    """Calculates the sampled log density of a Gaussian.

    Version for Tensor arrays as inputs.

    Parameters:
        sample:
            Sample at which to compute the density.
        mean:
            Mean of Gaussian distribution.
        logvar:
            Log variance of Gaussian distribution.
        all_combinations:
            Whether to return values for all combinations of batch pairs of `sample` and `mean`.
        axis:
            The dimension over which all combinations are computed. Default is 0 meaning the batch dimension.

    Returns:
        Tensor

    """
    if all_combinations:
        sample = tf.expand_dims(sample, axis=axis + 1)
        mean = tf.expand_dims(mean, axis=axis)
        logvar = tf.expand_dims(logvar, axis=axis)

    log_density = LOG2PI + logvar + tf.exp(-logvar) * (sample - mean)**2
    log_density *= -0.5
    return log_density

VAE.utils.math.reduce_logmeanexp 

reduce_logmeanexp(input_tensor, axis=None, keepdims=False)

Calculates Log(Mean(Exp(x))).

Version for Tensor arrays as inputs.

Parameters:

  • input_tensor (Tensor) –

    Input tensor.

  • axis (int, default: None ) –

    The dimensions to reduce.

  • keepdims (bool, default: False ) –

    Whether to keep the axis as singleton dimensions.

Returns:

  • Tensor –

    Tensor.

Source code in VAE/utils/math.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
def reduce_logmeanexp(input_tensor: tf.Tensor, axis: int = None, keepdims: bool = False) -> tf.Tensor:
    """Calculates Log(Mean(Exp(x))).

    Version for Tensor arrays as inputs.

    Parameters:
        input_tensor:
            Input tensor.
        axis:
            The dimensions to reduce.
        keepdims:
            Whether to keep the axis as singleton dimensions.

    Returns:
        Tensor.

    """
    lse = tf.reduce_logsumexp(input_tensor, axis=axis, keepdims=keepdims)
    n = tf.size(input_tensor) // tf.size(lse)
    log_n = tf.math.log(tf.cast(n, lse.dtype))
    return lse - log_n

VAE.utils.math.similarity 

similarity(input_tensor, temperature=1.0)

Returns the cosine similarity.

The function returns a similarity measure between all samples in a batch. First, a correlation matrix of the normalized input is obtained (the cosine similarity matrix). Then, the similarity measure is defined as the categorical cross entropy of the cosine similarity matrix of shape (batch_size, batch_size) with the identity matrix as target. The correlation can be scaled by temperature prior to the calculation of the cross entropy.

Parameters:

  • input_tensor (Tensor) –

    Input tensor of shape (batch_size, channels).

  • temperature (float, default: 1.0 ) –

    Temperature for the softmax.

Returns:

  • Tensor –

    Tensor of shape (batch_size,).

Source code in VAE/utils/math.py 
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
def similarity(input_tensor: tf.Tensor, temperature: float = 1.) -> tf.Tensor:
    """Returns the cosine similarity.

    The function returns a similarity measure between all samples in a batch. First, a correlation matrix of the
    normalized input is obtained (the cosine similarity matrix). Then, the similarity measure is defined as the
    categorical cross entropy of the cosine similarity matrix of shape `(batch_size, batch_size)` with the identity
    matrix as target. The correlation can be scaled by `temperature` prior to the calculation of the cross entropy.

    Parameters:
        input_tensor:
            Input tensor of shape `(batch_size, channels)`.
        temperature:
            Temperature for the softmax.

    Returns:
        Tensor of shape `(batch_size,)`.

    """
    # normalize input
    l2 = tf.math.l2_normalize(input_tensor, axis=-1)
    # correlation matrix
    similarity = tf.matmul(l2, l2, transpose_b=True)
    # apply temperature
    similarity /= temperature
    # target labels = diagonal elements
    labels = tf.range(tf.shape(similarity)[0])
    # cross entropy loss
    loss = ks.losses.sparse_categorical_crossentropy(labels, similarity, from_logits=True, axis=-1)
    return loss

VAE.utils.math.log_density_gaussian_numpy 

log_density_gaussian_numpy(sample, mean, logvar, all_combinations=False)

Calculates the sampled log density of a Gaussian.

Numpy version of :func:log_density_gaussian.

Parameters:

  • sample (ndarray) –

    Array of shape (batch_size, laten_dim) at which to compute the density.

  • mean (ndarray) –

    Array of shape (batch_size, laten_dim) represnting the mean of Gaussian distribution.

  • logvar (ndarray) –

    Array of shape (batch_size, laten_dim) represnting the log variance of Gaussian distribution.

  • all_combinations (bool, default: False ) –

    Whether to return values for all combinations of batch pairs of sample and mean.

Returns:

  • ndarray –

    If all_combinations=True returns an array of shape (batch_size, batch_size, latent_dim).

  • ndarray –

    If all_combinations=False returns an array of shape (batch_size, latent_dim).

Source code in VAE/utils/math.py 
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
def log_density_gaussian_numpy(sample: np.ndarray,
                               mean: np.ndarray,
                               logvar: np.ndarray,
                               all_combinations: bool = False) -> np.ndarray:
    """Calculates the sampled log density of a Gaussian.

    Numpy version of :func:`log_density_gaussian`.

    Parameters:
        sample:
            Array of shape `(batch_size, laten_dim)` at which to compute the density.
        mean:
            Array of shape `(batch_size, laten_dim)` represnting the mean of Gaussian distribution.
        logvar:
            Array of shape `(batch_size, laten_dim)` represnting the log variance of Gaussian distribution.
        all_combinations:
            Whether to return values for all combinations of batch pairs of `sample` and `mean`.

    Returns:
        If `all_combinations=True` returns an array of shape `(batch_size, batch_size, latent_dim)`.
        If `all_combinations=False` returns an array of shape `(batch_size, latent_dim)`.

    """
    if all_combinations:
        # reshape to (batch_size, 1, latent_dim)
        sample = np.expand_dims(sample, axis=1)

        # reshape to (1, batch_size, latent_dim)
        mean = np.expand_dims(mean, axis=0)
        logvar = np.expand_dims(logvar, axis=0)

    log_density = np.log(2 * np.pi) + logvar + np.exp(-logvar) * (sample - mean)**2
    log_density *= -0.5
    return log_density

VAE.utils.math.logmeanexp_numpy 

logmeanexp_numpy(x, axis=None, keepdims=False)

Calculates Log(Mean(Exp(x))).

Numpy version of :func:logmeanexp.

Parameters:

  • x (ndarray) –

    Input array.

  • axis (int, default: None ) –

    The dimensions to reduce.

  • keepdims (bool, default: False ) –

    Whether to keep the axis as singleton dimensions.

Returns:

  • ndarray –

    Numpy array.

Source code in VAE/utils/math.py 
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
def logmeanexp_numpy(x: np.ndarray, axis: int = None, keepdims: bool = False) -> np.ndarray:
    """Calculates Log(Mean(Exp(x))).

    Numpy version of :func:`logmeanexp`.

    Parameters:
        x:
            Input array.
        axis:
            The dimensions to reduce.
        keepdims:
            Whether to keep the axis as singleton dimensions.

    Returns:
        Numpy array.

    """
    lse = logsumexp(x, axis=axis, keepdims=keepdims)
    n = np.size(x) // np.size(lse)
    log_n = np.log(n)
    return lse - log_n

VAE.utils.math.kl_divergence_numpy 

kl_divergence_numpy(mean1, logvar1, mean2=0, logvar2=0, all_combinations=False)

Calculates the KL divergence between two multivariate Gaussian distributions with diagonal covariance matrices.

The KL divergence is returned for each latent dimensions separately. Since the KL divergence is additive in the case of diagonal covariance matrices, the KL divergence can be obtained from the sum over the last dimensions of the output.

Parameters:

  • mean1 (float) –

    Array of shape (batch_size, laten_dim) representing the mean of the first Gaussian distribution.

  • logvar1 (float) –

    Array of shape (batch_size, laten_dim) representing the log variance of the second Gaussian distribution.

  • mean2 (float, default: 0 ) –

    Array of shape (batch_size, laten_dim) representing the mean of the second Gaussian distribution. Defaults to 0.

  • logvar2 (float, default: 0 ) –

    Array of shape (batch_size, laten_dim) representing the log variance of the second Gaussian distribution. Defaults to 0.

  • all_combinations (bool, default: False ) –

    Whether to return values for all combinations of batch pairs of sample and mean.

Returns:

  • ndarray –

    If all_combinations=True returns an array of shape (batch_size, batch_size, latent_dim).

  • ndarray –

    If all_combinations=False returns an array of shape (batch_size, latent_dim).

Source code in VAE/utils/math.py 
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
def kl_divergence_numpy(mean1: float,
                        logvar1: float,
                        mean2: float = 0,
                        logvar2: float = 0,
                        all_combinations: bool = False) -> np.ndarray:
    """Calculates the KL divergence between two multivariate Gaussian distributions with diagonal covariance matrices.

    The KL divergence is returned for each latent dimensions separately. Since the KL divergence is additive in the case
    of diagonal covariance matrices, the KL divergence can be obtained from the sum over the last dimensions of the
    output.

    Parameters:
        mean1:
            Array of shape `(batch_size, laten_dim)` representing the mean of the first Gaussian distribution.
        logvar1:
            Array of shape `(batch_size, laten_dim)` representing the log variance of the second Gaussian distribution.
        mean2:
            Array of shape `(batch_size, laten_dim)` representing the mean of the second Gaussian distribution. Defaults
            to 0.
        logvar2:
            Array of shape `(batch_size, laten_dim)` representing the log variance of the second Gaussian distribution.
            Defaults to 0.
        all_combinations:
            Whether to return values for all combinations of batch pairs of `sample` and `mean`.

    Returns:
        If `all_combinations=True` returns an array of shape `(batch_size, batch_size, latent_dim)`.
        If `all_combinations=False` returns an array of shape `(batch_size, latent_dim)`.

    """
    if all_combinations:
        # reshape to (batch_size, 1, latent_dim)
        mean1 = np.expand_dims(mean1, axis=1)
        logvar1 = np.expand_dims(logvar1, axis=1)

        # reshape to (1, batch_size, latent_dim)
        mean2 = np.expand_dims(mean2, axis=0)
        logvar2 = np.expand_dims(logvar2, axis=0)

    kl_div = logvar2 - logvar1 + (np.exp(logvar1) + (mean1 - mean2)**2) * np.exp(-logvar2) - 1
    kl_div *= 0.5
    return kl_div

VAE.utils.math.wasserstein_distance_numpy 

wasserstein_distance_numpy(mean1, logvar1, mean2=0, logvar2=0, all_combinations=False)

Calculates the Wasserstein distance between two multivariate Gaussian distributions with diagonal covariance matrices.

The Wasserstein distance is returned for each latent dimensions separately. Since the Wasserstein distance is additive in the case of diagonal covariance matrices, the Wasserstein distance can be obtained from the sum over the last dimensions of the output.

Parameters:

  • mean1 (float) –

    Array of shape (batch_size, laten_dim) representing the mean of the first Gaussian distribution.

  • logvar1 (float) –

    Array of shape (batch_size, laten_dim) representing the log variance of the second Gaussian distribution.

  • mean2 (float, default: 0 ) –

    Array of shape (batch_size, laten_dim) representing the mean of the second Gaussian distribution. Defaults to 0.

  • logvar2 (float, default: 0 ) –

    Array of shape (batch_size, laten_dim) representing the log variance of the second Gaussian distribution. Defaults to 0.

  • all_combinations (bool, default: False ) –

    Whether to return values for all combinations of batch pairs of sample and mean.

Returns:

  • ndarray –

    If all_combinations=True returns an array of shape (batch_size, batch_size, latent_dim).

  • ndarray –

    If all_combinations=False returns an array of shape (batch_size, latent_dim).

Source code in VAE/utils/math.py 
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
def wasserstein_distance_numpy(mean1: float,
                               logvar1: float,
                               mean2: float = 0,
                               logvar2: float = 0,
                               all_combinations: bool = False) -> np.ndarray:
    """Calculates the Wasserstein distance between two multivariate Gaussian distributions with diagonal covariance
    matrices.

    The Wasserstein distance is returned for each latent dimensions separately. Since the Wasserstein distance is
    additive in the case of diagonal covariance matrices, the Wasserstein distance can be obtained from the sum over the
    last dimensions of the output.

    Parameters:
        mean1:
            Array of shape `(batch_size, laten_dim)` representing the mean of the first Gaussian distribution.
        logvar1:
            Array of shape `(batch_size, laten_dim)` representing the log variance of the second Gaussian distribution.
        mean2:
            Array of shape `(batch_size, laten_dim)` representing the mean of the second Gaussian distribution. Defaults
            to 0.
        logvar2:
            Array of shape `(batch_size, laten_dim)` representing the log variance of the second Gaussian distribution.
            Defaults to 0.
        all_combinations:
            Whether to return values for all combinations of batch pairs of `sample` and `mean`.

    Returns:
        If `all_combinations=True` returns an array of shape `(batch_size, batch_size, latent_dim)`.
        If `all_combinations=False` returns an array of shape `(batch_size, latent_dim)`.

    """
    if all_combinations:
        # reshape to (batch_size, 1, latent_dim)
        mean1 = np.expand_dims(mean1, axis=1)
        logvar1 = np.expand_dims(logvar1, axis=1)

        # reshape to (1, batch_size, latent_dim)
        mean2 = np.expand_dims(mean2, axis=0)
        logvar2 = np.expand_dims(logvar2, axis=0)

    wd = (mean1 - mean2)**2 + np.exp(logvar1) + np.exp(logvar2) - 2 * np.exp(0.5 * (logvar1 + logvar2))

    return wd