def test_mixture_logprob(
distribution: Distribution,
values_outside_support: Tensor,
distribution_output: DistributionOutput,
) -> None:
assert np.all(
~np.isnan(distribution.log_prob(values_outside_support).asnumpy())
), f"{distribution} should return -inf log_probs instead of NaNs"
p = 0.5
gaussian = Gaussian(mu=mx.nd.array([0]), sigma=mx.nd.array([2.0]))
mixture = MixtureDistribution(
mixture_probs=mx.nd.array([[p, 1 - p]]),
components=[gaussian, distribution],
)
lp = mixture.log_prob(values_outside_support)
assert np.allclose(
lp.asnumpy(),
np.log(p) + gaussian.log_prob(values_outside_support).asnumpy(),
atol=1e-6,
), f"log_prob(x) should be equal to log(p)+gaussian.log_prob(x)"
fit_mixture = fit_mixture_distribution(
values_outside_support,
MixtureDistributionOutput([GaussianOutput(), distribution_output]),
variate_dimensionality=1,
epochs=3,
)
for ci, c in enumerate(fit_mixture.components):
for ai, a in enumerate(c.args):
assert ~np.isnan(a.asnumpy()), f"NaN gradients led to {c}"
def test_deterministic_l2(mu: float, hybridize: bool) -> None:
"""
Test to check that maximizing the likelihood recovers the parameters.
This tests uses the Gaussian distribution with fixed variance and sample mean.
This essentially reduces to determistic L2.
"""
# generate samples
mu = mu
mus = mx.nd.zeros(NUM_SAMPLES) + mu
deterministic_distr = Gaussian(mu=mus, sigma=0.1 * mx.nd.ones_like(mus))
samples = deterministic_distr.sample()
class GaussianFixedVarianceOutput(GaussianOutput):
@classmethod
def domain_map(cls, F, mu, sigma):
sigma = 0.1 * F.ones_like(sigma)
return mu.squeeze(axis=-1), sigma.squeeze(axis=-1)
mu_hat, _ = maximum_likelihood_estimate_sgd(
GaussianFixedVarianceOutput(),
samples,
init_biases=[3 * mu, 0.1],
hybridize=hybridize,
num_epochs=PositiveInt(1),
)
assert (np.abs(mu_hat - mu) <
TOL * mu), f"mu did not match: mu = {mu}, mu_hat = {mu_hat}"
def test_gaussian_likelihood(mu: float, sigma: float, hybridize: bool):
"""
Test to check that maximizing the likelihood recovers the parameters
"""
# generate samples
mus = mx.nd.zeros((NUM_SAMPLES, )) + mu
sigmas = mx.nd.zeros((NUM_SAMPLES, )) + sigma
distr = Gaussian(mus, sigmas)
samples = distr.sample()
init_biases = [
mu - START_TOL_MULTIPLE * TOL * mu,
inv_softplus(sigma - START_TOL_MULTIPLE * TOL * sigma),
]
mu_hat, sigma_hat = maximum_likelihood_estimate_sgd(
GaussianOutput(),
samples,
init_biases=init_biases,
hybridize=hybridize,
learning_rate=PositiveFloat(0.001),
num_epochs=PositiveInt(5),
)
assert (np.abs(mu_hat - mu) <
TOL * mu), f"mu did not match: mu = {mu}, mu_hat = {mu_hat}"
assert (np.abs(sigma_hat - sigma) < TOL * sigma
), f"alpha did not match: sigma = {sigma}, sigma_hat = {sigma_hat}"
def test_box_cox_tranform(
lam_1: float, lam_2: float, mu: float, sigma: float, hybridize: bool
):
"""
Test to check that maximizing the likelihood recovers the parameters
"""
# generate samples
lamdas_1 = mx.nd.zeros((NUM_SAMPLES,)) + lam_1
lamdas_2 = mx.nd.zeros((NUM_SAMPLES,)) + lam_2
transform = InverseBoxCoxTransform(lamdas_1, lamdas_2)
mus = mx.nd.zeros((NUM_SAMPLES,)) + mu
sigmas = mx.nd.zeros((NUM_SAMPLES,)) + sigma
gausian_distr = Gaussian(mus, sigmas)
# Here the base distribution is Guassian which is transformed to
# non-Gaussian via the inverse Box-Cox transform.
# Sampling from `trans_distr` gives non-Gaussian samples
trans_distr = TransformedDistribution(gausian_distr, [transform])
# Given the non-Gaussian samples find the true parameters
# of the Box-Cox transformation as well as the underlying Gaussian distribution.
samples = trans_distr.sample()
init_biases = [
mu - START_TOL_MULTIPLE * TOL * mu,
inv_softplus(sigma - START_TOL_MULTIPLE * TOL * sigma),
lam_1 - START_TOL_MULTIPLE * TOL * lam_1,
inv_softplus(lam_2 - START_TOL_MULTIPLE * TOL * lam_2),
]
mu_hat, sigma_hat, lam_1_hat, lam_2_hat = maximum_likelihood_estimate_sgd(
TransformedDistributionOutput(
GaussianOutput(),
InverseBoxCoxTransformOutput(lb_obs=lam_2, fix_lambda_2=True),
),
samples,
init_biases=init_biases,
hybridize=hybridize,
learning_rate=PositiveFloat(0.01),
num_epochs=PositiveInt(18),
)
assert (
np.abs(lam_1_hat - lam_1) < TOL * lam_1
), f"lam_1 did not match: lam_1 = {lam_1}, lam_1_hat = {lam_1_hat}"
# assert (
# np.abs(lam_2_hat - lam_2) < TOL * lam_2
# ), f"lam_2 did not match: lam_2 = {lam_2}, lam_2_hat = {lam_2_hat}"
assert np.abs(mu_hat - mu) < TOL * np.abs(
mu
), f"mu did not match: mu = {mu}, mu_hat = {mu_hat}"
assert (
np.abs(sigma_hat - sigma) < TOL * sigma
), f"sigma did not match: sigma = {sigma}, sigma_hat = {sigma_hat}"
NUM_SAMPLES = 1_000
NUM_SAMPLES_LARGE = 1_000_000
SHAPE = (2, 1, 3)
np.random.seed(35120171)
mx.random.seed(35120171)
@pytest.mark.parametrize(
"distr1, distr2, p",
[
(
Gaussian(
mu=mx.nd.zeros(shape=SHAPE),
sigma=1e-3 + 0.2 * mx.nd.ones(shape=SHAPE),
),
Gaussian(
mu=mx.nd.ones(shape=SHAPE),
sigma=1e-3 + 0.1 * mx.nd.ones(shape=SHAPE),
),
0.2 * mx.nd.ones(shape=SHAPE),
),
(
StudentT(
mu=mx.nd.ones(shape=SHAPE),
sigma=1e-1 + mx.nd.zeros(shape=SHAPE),
nu=mx.nd.zeros(shape=SHAPE) + 2.2,
),
Gaussian(
mu=-mx.nd.ones(shape=SHAPE),
NUM_SAMPLES = 1_000
NUM_SAMPLES_LARGE = 1_000_000
SHAPE = (2, 1, 3)
np.random.seed(35120171)
mx.random.seed(35120171)
@pytest.mark.parametrize(
"distr1, distr2, p",
[
(
Gaussian(
mu=mx.nd.zeros(shape=SHAPE),
sigma=1e-3 + 0.2 * mx.nd.ones(shape=SHAPE),
),
Gaussian(
mu=mx.nd.ones(shape=SHAPE),
sigma=1e-3 + 0.1 * mx.nd.ones(shape=SHAPE),
),
0.2 * mx.nd.ones(shape=SHAPE),
),
(
StudentT(
mu=mx.nd.ones(shape=SHAPE),
sigma=1e-1 + mx.nd.zeros(shape=SHAPE),
nu=mx.nd.zeros(shape=SHAPE) + 2.2,
),
Gaussian(
mu=-mx.nd.ones(shape=SHAPE),
Poisson,
StudentT,
TransformedDistribution,
Uniform,
)
from gluonts.mx.distribution.bijection import AffineTransformation
from gluonts.mx.distribution.box_cox_transform import BoxCoxTransform
from gluonts.mx.util import make_nd_diag
@pytest.mark.parametrize(
"distr, expected_batch_shape, expected_event_shape",
[
(
Gaussian(
mu=mx.nd.zeros(shape=(3, 4, 5)),
sigma=mx.nd.ones(shape=(3, 4, 5)),
),
(3, 4, 5),
(),
),
(
Gamma(
alpha=mx.nd.ones(shape=(3, 4, 5)),
beta=mx.nd.ones(shape=(3, 4, 5)),
),
(3, 4, 5),
(),
),
(
Beta(
alpha=mx.nd.ones(shape=(3, 4, 5)),
