def sample(self, mean: Tensor, covariance: Tensor) -> Tensor:
r"""
Parameters
----------
covariance
The covariance matrix of the GP of shape (batch_size, prediction_length, prediction_length).
mean
The mean vector of the GP of shape (batch_size, prediction_length).
Returns
-------
Tensor
Samples from a Gaussian Process of shape (batch_size, prediction_length, num_samples), where :math:`L`
is the matrix square root, Cholesky Factor of the covariance matrix with the added noise tolerance on the
diagonal, :math:`Lz`, where :math:`z \sim N(0,I)` and assumes the mean is zero.
"""
assert (self.num_samples
is not None), "The value of `num_samples` must be set."
assert (self.prediction_length
is not None), "The value of `prediction_length` must be set."
samples = MultivariateGaussian(
mean,
self._compute_cholesky_gp(covariance, self.prediction_length,
self.sample_noise),
).sample_rep(self.num_samples, dtype=self.float_type
) # Shape (num_samples, batch_size, prediction_length)
return self.F.transpose(samples, axes=(1, 2, 0))
def log_prob(self, x_train: Tensor, y_train: Tensor) -> Tensor:
r"""
This method computes the negative marginal log likelihood
.. math::
:nowrap:
\begin{aligned}
\frac{1}{2} [d \log(2\pi) + \log(|K|) + y^TK^{-1}y],
\end{aligned}
where :math:`d` is the number of data points.
This can be written in terms of the Cholesky factor :math:`L` as
.. math::
:nowrap:
\begin{aligned}
\log(|K|) = \log(|LL^T|) &= \log(|L||L|^T) = \log(|L|^2) = 2\log(|L|) \\
&= 2\log\big(\prod_i^n L_{ii}\big) = 2 \sum_i^N \log(L_{ii})
\end{aligned}
and
.. math::
:nowrap:
\begin{aligned}
y^TK^{-1}y = (y^TL^{-T})(L^{-1}y) = (L^{-1}y)^T(L^{-1}y) = ||L^{-1}y||_2^2.
\end{aligned}
Parameters
--------------------
x_train
Training set of features of shape (batch_size, context_length, num_features).
y_train
Training labels of shape (batch_size, context_length).
Returns
--------------------
Tensor
The negative log marginal likelihood of shape (batch_size,)
"""
assert (
self.context_length is not None
), "The value of `context_length` must be set."
return -MultivariateGaussian(
self.F.zeros_like(y_train), # 0 mean gaussian process prior
self._compute_cholesky_gp(
self.kernel.kernel_matrix(x_train, x_train),
self.context_length,
),
).log_prob(y_train)
def test_multivariate_gaussian(hybridize: bool) -> None:
num_samples = 2000
dim = 2
mu = np.arange(0, dim) / float(dim)
L_diag = np.ones((dim, ))
L_low = 0.1 * np.ones((dim, dim)) * np.tri(dim, k=-1)
L = np.diag(L_diag) + L_low
Sigma = L.dot(L.transpose())
distr = MultivariateGaussian(mu=mx.nd.array(mu), L=mx.nd.array(L))
samples = distr.sample(num_samples)
mu_hat, L_hat = maximum_likelihood_estimate_sgd(
MultivariateGaussianOutput(dim=dim),
samples,
init_biases=
None, # todo we would need to rework biases a bit to use it in the multivariate case
hybridize=hybridize,
learning_rate=PositiveFloat(0.01),
num_epochs=PositiveInt(10),
)
distr = MultivariateGaussian(mu=mx.nd.array([mu_hat]),
L=mx.nd.array([L_hat]))
Sigma_hat = distr.variance[0].asnumpy()
assert np.allclose(
mu_hat, mu, atol=0.1,
rtol=0.1), f"mu did not match: mu = {mu}, mu_hat = {mu_hat}"
assert np.allclose(
Sigma_hat, Sigma, atol=0.1, rtol=0.1
), f"Sigma did not match: sigma = {Sigma}, sigma_hat = {Sigma_hat}"
),
(3, 4, 5),
(),
),
(
StudentT(
mu=mx.nd.zeros(shape=(3, 4, 5)),
sigma=mx.nd.ones(shape=(3, 4, 5)),
nu=mx.nd.ones(shape=(3, 4, 5)),
),
(3, 4, 5),
(),
),
(
MultivariateGaussian(
mu=mx.nd.zeros(shape=(3, 4, 5)),
L=make_nd_diag(F=mx.nd, x=mx.nd.ones(shape=(3, 4, 5)), d=5),
),
(3, 4),
(5, ),
),
(Dirichlet(alpha=mx.nd.ones(shape=(3, 4, 5))), (3, 4), (5, )),
(
DirichletMultinomial(
dim=5, n_trials=9, alpha=mx.nd.ones(shape=(3, 4, 5))),
(3, 4),
(5, ),
),
(
Laplace(mu=mx.nd.zeros(shape=(3, 4, 5)),
b=mx.nd.ones(shape=(3, 4, 5))),
(3, 4, 5),