def test_lowrank_multivariate_gaussian(hybridize: bool) -> None:
num_samples = 2000
dim = 2
rank = 1
mu = np.arange(0, dim) / float(dim)
D = np.eye(dim) * (np.arange(dim) / dim + 0.5)
W = np.sqrt(np.ones((dim, rank)) * 0.2)
Sigma = D + W.dot(W.transpose())
distr = LowrankMultivariateGaussian(
mu=mx.nd.array([mu]),
D=mx.nd.array([np.diag(D)]),
W=mx.nd.array([W]),
dim=dim,
rank=rank,
)
assert np.allclose(
distr.variance[0].asnumpy(), Sigma, atol=0.1, rtol=0.1
), f"did not match: sigma = {Sigma}, sigma_hat = {distr.variance[0]}"
samples = distr.sample(num_samples).squeeze().asnumpy()
mu_hat, D_hat, W_hat = maximum_likelihood_estimate_sgd(
LowrankMultivariateGaussianOutput(dim=dim,
rank=rank,
sigma_init=0.2,
sigma_minimum=0.0),
samples,
learning_rate=PositiveFloat(0.01),
num_epochs=PositiveInt(25),
init_biases=
None, # todo we would need to rework biases a bit to use it in the multivariate case
hybridize=hybridize,
)
distr = LowrankMultivariateGaussian(
dim=dim,
rank=rank,
mu=mx.nd.array([mu_hat]),
D=mx.nd.array([D_hat]),
W=mx.nd.array([W_hat]),
)
Sigma_hat = distr.variance.asnumpy()
assert np.allclose(
mu_hat, mu, atol=0.2,
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}"
def loss(self, F, target: Tensor, distr: Distribution) -> Tensor:
"""
Computes loss given the output of the network in the form of
distribution. The loss is given by:
`self.CRPS_weight` * `loss_CRPS` + `self.likelihood_weight` *
`neg_likelihoods`,
where
* `loss_CRPS` is computed on the samples drawn from the predicted
`distr` (optionally after reconciling them),
* `neg_likelihoods` are either computed directly using the predicted
`distr` or from the estimated distribution based on (coherent)
samples, depending on the `sample_LH` flag.
Parameters
----------
F
target
Tensor with shape (batch_size, seq_len, target_dim)
distr
Distribution instances
Returns
-------
Loss
Tensor with shape (batch_size, seq_length, 1)
"""
# Sample from the predicted distribution if we are computing CRPS loss
# or likelihood using the distribution based on (coherent) samples.
# Samples shape: (num_samples, batch_size, seq_len, target_dim)
if self.sample_LH or (self.CRPS_weight > 0.0):
samples = self.get_samples_for_loss(distr=distr)
if self.sample_LH:
# Estimate the distribution based on (coherent) samples.
distr = LowrankMultivariateGaussian.fit(F, samples=samples, rank=0)
neg_likelihoods = -distr.log_prob(target).expand_dims(axis=-1)
loss_CRPS = F.zeros_like(neg_likelihoods)
if self.CRPS_weight > 0.0:
loss_CRPS = (
EmpiricalDistribution(samples=samples, event_dim=1)
.crps_univariate(x=target)
.expand_dims(axis=-1)
)
return (
self.CRPS_weight * loss_CRPS
+ self.likelihood_weight * neg_likelihoods
)
def test_gpvar_proj():
# test that gp proj gives expected shapes
batch = 1
hidden_size = 3
dim = 4
rank = 2
states = mx.ndarray.ones(shape=(batch, dim, hidden_size))
gp_proj = GPArgProj(rank=rank)
gp_proj.initialize()
distr_args = gp_proj(states)
mu, D, W = distr_args
assert mu.shape == (batch, dim)
assert D.shape == (batch, dim)
assert W.shape == (batch, dim, rank)
distr = LowrankMultivariateGaussian(dim, rank, *distr_args)
assert distr.mean.shape == (batch, dim)
def test_empirical_distribution(hybridize: bool) -> None:
r"""
This verifies if the loss implemented by `EmpiricalDistribution` is correct.
This is done by recovering parameters of a parametric distribution not by maximizing likelihood but by
optimizing CRPS loss on the Monte Carlo samples drawn from the underlying parametric distribution.
More precisely, given observations `obs` drawn from the true distribution p(x; \theta^*), we solve
\theta_hat = \argmax_{\theta} CRPS(obs, {x}_i)
subject to: x_i ~ p(x; \theta)
and verify if \theta^* and \theta_hat agree.
This test uses Multivariate Gaussian with diagonal covariance. Once multivariate CRPS is implemented in
`EmpiricalDistribution` one could use `LowrankMultivariateGaussian` as well. Any univariate distribution whose
`sample_rep` is differentiable can also be used in this test.
"""
num_obs = 2000
dim = 2
# Multivariate CRPS is not implemented in `EmpiricalDistribution`.
rank = 0
W = None
mu = np.arange(0, dim) / float(dim)
D = np.eye(dim) * (np.arange(dim) / dim + 0.5)
Sigma = D
distr = LowrankMultivariateGaussian(
mu=mx.nd.array([mu]),
D=mx.nd.array([np.diag(D)]),
W=W,
dim=dim,
rank=rank,
)
obs = distr.sample(num_obs).squeeze().asnumpy()
theta_hat = maximum_likelihood_estimate_sgd(
EmpiricalDistributionOutput(
num_samples=200,
distr_output=LowrankMultivariateGaussianOutput(dim=dim,
rank=rank,
sigma_init=0.2,
sigma_minimum=0.0),
),
obs,
learning_rate=PositiveFloat(0.01),
num_epochs=PositiveInt(25),
init_biases=
None, # todo we would need to rework biases a bit to use it in the multivariate case
hybridize=hybridize,
)
mu_hat, D_hat = theta_hat
W_hat = None
distr = LowrankMultivariateGaussian(
dim=dim,
rank=rank,
mu=mx.nd.array([mu_hat]),
D=mx.nd.array([D_hat]),
W=W_hat,
)
Sigma_hat = distr.variance.asnumpy()
print(mu_hat, Sigma_hat)
assert np.allclose(
mu_hat, mu, atol=0.2,
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}"