def test_studentT_likelihood(
mu: float, sigma: float, nu: float, hybridize: bool
) -> None:
"""
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
nus = mx.nd.zeros((NUM_SAMPLES,)) + nu
distr = StudentT(mus, sigmas, nus)
samples = distr.sample()
# nu takes very long to learn, so we initialize it at the true value.
# transform used is softplus(x) + 2
init_bias = [
mu - START_TOL_MULTIPLE * TOL * mu,
inv_softplus(sigma - START_TOL_MULTIPLE * TOL * sigma),
inv_softplus(nu - 2),
]
mu_hat, sigma_hat, nu_hat = maximum_likelihood_estimate_sgd(
StudentTOutput(),
samples,
init_biases=init_bias,
hybridize=hybridize,
num_epochs=PositiveInt(10),
learning_rate=PositiveFloat(1e-2),
)
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"sigma did not match: sigma = {sigma}, sigma_hat = {sigma_hat}"
assert (
np.abs(nu_hat - nu) < TOL * nu
), "nu0 did not match: nu0 = %s, nu_hat = %s" % (nu, nu_hat)
[
(
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),
sigma=1e-1 + mx.nd.zeros(shape=SHAPE),
),
mx.nd.random_uniform(shape=SHAPE),
),
# TODO: add a multivariate case here
],
)
@pytest.mark.parametrize("serialize_fn", serialize_fn_list)
def test_mixture(
distr1: Distribution, distr2: Distribution, p: Tensor, serialize_fn
) -> None:
),
Gaussian(
mu=mx.nd.zeros(shape=BATCH_SHAPE),
sigma=mx.nd.ones(shape=BATCH_SHAPE),
),
Gamma(
alpha=mx.nd.ones(shape=BATCH_SHAPE),
beta=mx.nd.ones(shape=BATCH_SHAPE),
),
Beta(
alpha=0.5 * mx.nd.ones(shape=BATCH_SHAPE),
beta=0.5 * mx.nd.ones(shape=BATCH_SHAPE),
),
StudentT(
mu=mx.nd.zeros(shape=BATCH_SHAPE),
sigma=mx.nd.ones(shape=BATCH_SHAPE),
nu=mx.nd.ones(shape=BATCH_SHAPE),
),
Dirichlet(alpha=mx.nd.ones(shape=BATCH_SHAPE)),
Laplace(mu=mx.nd.zeros(shape=BATCH_SHAPE),
b=mx.nd.ones(shape=BATCH_SHAPE)),
NegativeBinomial(
mu=mx.nd.zeros(shape=BATCH_SHAPE),
alpha=mx.nd.ones(shape=BATCH_SHAPE),
),
Poisson(rate=mx.nd.ones(shape=BATCH_SHAPE)),
Uniform(
low=-mx.nd.ones(shape=BATCH_SHAPE),
high=mx.nd.ones(shape=BATCH_SHAPE),
),
PiecewiseLinear(
),
(3, 4, 5),
(),
),
(
Beta(
alpha=mx.nd.ones(shape=(3, 4, 5)),
beta=mx.nd.ones(shape=(3, 4, 5)),
),
(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, )),
(