def test_multivariate_normal() -> 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 = MultivariateNormal(loc=torch.Tensor(mu), scale_tril=torch.Tensor(L)) samples = distr.sample((num_samples,)) mu_hat, L_hat = maximum_likelihood_estimate_sgd( MultivariateNormalOutput(dim=dim), samples, init_biases=None, # todo we would need to rework biases a bit to use it in the multivariate case learning_rate=0.01, num_epochs=10, ) distr = MultivariateNormal(loc=torch.tensor(mu_hat), scale_tril=torch.tensor(L_hat)) Sigma_hat = distr.covariance_matrix.numpy() 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}"
), ( LowRankMultivariateNormalOutput(dim=target_dim, rank=2), 10, estimator, True, ), ( LowRankMultivariateNormalOutput(dim=target_dim, rank=2), 10, estimator, False, ), (None, 10, estimator, True), ( MultivariateNormalOutput(dim=target_dim), 10, estimator, True, ), ( MultivariateNormalOutput(dim=target_dim), 10, estimator, False, ), ], ) def test_deepvar( distr_output, num_batches_per_epoch, Estimator, use_marginal_transformation, ):
# @pytest.mark.timeout(10) @pytest.mark.parametrize( "distr_output, num_batches_per_epoch, Estimator, use_marginal_transformation", [ (NormalOutput(dim=target_dim), 10, estimator, True,), (NormalOutput(dim=target_dim), 10, estimator, False,), (LowRankMultivariateNormalOutput(dim=target_dim, rank=2), 10, estimator, True,), ( LowRankMultivariateNormalOutput(dim=target_dim, rank=2), 10, estimator, False, ), (None, 10, estimator, True), (MultivariateNormalOutput(dim=target_dim), 10, estimator, True,), (MultivariateNormalOutput(dim=target_dim), 10, estimator, False,), ], ) def test_deepvar( distr_output, num_batches_per_epoch, Estimator, use_marginal_transformation, ): estimator = Estimator( input_size=47, num_cells=20, num_layers=1, dropout_rate=0.0, pick_incomplete=True, target_dim=target_dim, prediction_length=metadata.prediction_length,