Source code for

# Copyright 2018, Inc. or its affiliates. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 (the "License").
# You may not use this file except in compliance with the License.
# A copy of the License is located at
# or in the "license" file accompanying this file. This file is distributed
# express or implied. See the License for the specific language governing
# permissions and limitations under the License.

import math
from typing import Optional, Tuple

import numpy as np

from gluonts.core.component import validated
from import Tensor
from import lower_triangular_ones

from .distribution import Distribution, _sample_multiple, getF
from .distribution_output import DistributionOutput

[docs]class MultivariateGaussian(Distribution): r""" Multivariate Gaussian distribution, specified by the mean vector and the Cholesky factor of its covariance matrix. Parameters ---------- mu mean vector, of shape (..., d) L Lower triangular Cholesky factor of covariance matrix, of shape (..., d, d) F A module that can either refer to the Symbol API or the NDArray API in MXNet """ is_reparameterizable = True @validated() def __init__(self, mu: Tensor, L: Tensor, F=None) -> None: = mu self.L = L @property def F(self): return getF( def __getitem__(self, item): raise NotImplementedError() @property def batch_shape(self) -> Tuple: return[:-1] @property def event_shape(self) -> Tuple: return[-1:] @property def event_dim(self) -> int: return 1
[docs] def log_prob(self, x: Tensor) -> Tensor: # todo add an option to compute loss on diagonal covariance only to save time F = self.F # remark we compute d from the tensor but we could ask it to the user alternatively d = F.ones_like( residual = (x - # L^{-1} * (x - mu) L_inv_times_residual = F.linalg_trsm(self.L, residual) ll = ( F.broadcast_sub( -d / 2 * math.log(2 * math.pi), F.linalg_sumlogdiag(self.L) ) - 1 / 2 * F.linalg_syrk(L_inv_times_residual, transpose=True).squeeze() ) return ll
@property def mean(self) -> Tensor: return @property def variance(self) -> Tensor: return self.F.linalg_gemm2(self.L, self.L, transpose_b=True)
[docs] def sample_rep( self, num_samples: Optional[int] = None, dtype=np.float32 ) -> Tensor: r""" Draw samples from the multivariate Gaussian distributions. Internally, Cholesky factorization of the covariance matrix is used: sample = L v + mu, where L is the Cholesky factor, v is a standard normal sample. Parameters ---------- num_samples Number of samples to be drawn. dtype Data-type of the samples. Returns ------- Tensor Tensor with shape (num_samples, ..., d). """ def s(mu: Tensor, L: Tensor) -> Tensor: F = self.F samples_std_normal = F.sample_normal( mu=F.zeros_like(mu), sigma=F.ones_like(mu), dtype=dtype, ).expand_dims(axis=-1) samples = ( F.linalg_gemm2(L, samples_std_normal).squeeze(axis=-1) + mu ) return samples return _sample_multiple( s,, L=self.L, num_samples=num_samples )
[docs]class MultivariateGaussianOutput(DistributionOutput): @validated() def __init__(self, dim: int) -> None: self.args_dim = {"mu": dim, "Sigma": dim * dim} self.distr_cls = MultivariateGaussian self.dim = dim self.mask = None
[docs] def domain_map(self, F, mu_vector, L_vector): # apply softplus to the diagonal of L and mask upper coefficient to make it lower-triangular # diagonal matrix whose elements are diagonal elements of L mapped through a softplus d = self.dim # reshape from vector form (..., d * d) to matrix form(..., d, d) L_matrix = L_vector.reshape((-2, d, d, -4), reverse=1) L_diag = F.broadcast_mul( F.Activation( F.broadcast_mul(L_matrix, F.eye(d)), act_type="softrelu" ), F.eye(d), ) mask = lower_triangular_ones(F, d, offset=1) L_low = F.broadcast_mul(L_matrix, mask) return mu_vector, L_diag + L_low
@property def event_shape(self) -> Tuple: return (self.dim,)