diff --git a/pymc/distributions/continuous.py b/pymc/distributions/continuous.py index 0980356432..7f68c914d1 100644 --- a/pymc/distributions/continuous.py +++ b/pymc/distributions/continuous.py @@ -1599,6 +1599,13 @@ def dist(cls, b, kappa, mu=0, *args, **kwargs): return super().dist([b, kappa, mu], *args, **kwargs) + def get_moment(rv, size, b, kappa, mu): + mean = mu - (kappa - 1 / kappa) / b + + if not rv_size_is_none(size): + mean = at.full(size, mean) + return mean + def logp(value, b, kappa, mu): """ Calculate log-probability of Asymmetric-Laplace distribution at specified value. @@ -3012,6 +3019,12 @@ def dist(cls, alpha=1, mu=0.0, sigma=None, tau=None, sd=None, *args, **kwargs): return super().dist([mu, sigma, alpha], *args, **kwargs) + def get_moment(rv, size, mu, sigma, alpha): + mean = mu + sigma * (2 / np.pi) ** 0.5 * alpha / (1 + alpha ** 2) ** 0.5 + if not rv_size_is_none(size): + mean = at.full(size, mean) + return mean + def logp(value, mu, sigma, alpha): """ Calculate log-probability of SkewNormal distribution at specified value. diff --git a/pymc/tests/test_distributions_moments.py b/pymc/tests/test_distributions_moments.py index f52ca8baa0..86bc58a3eb 100644 --- a/pymc/tests/test_distributions_moments.py +++ b/pymc/tests/test_distributions_moments.py @@ -4,6 +4,7 @@ from scipy import special from pymc.distributions import ( + AsymmetricLaplace, Bernoulli, Beta, Binomial, @@ -29,6 +30,7 @@ Normal, Pareto, Poisson, + SkewNormal, StudentT, TruncatedNormal, Uniform, @@ -612,3 +614,59 @@ def test_discrete_uniform_moment(lower, upper, size, expected): with Model() as model: DiscreteUniform("x", lower=lower, upper=upper, size=size) assert_moment_is_expected(model, expected) + + +@pytest.mark.parametrize( + "alpha, mu, sigma, size, expected", + [ + (1.0, 1.0, 1.0, None, 1.56418958), + (1, np.ones(5), 1, None, np.full(5, 1.56418958)), + (np.ones(5), 1, np.ones(5), None, np.full(5, 1.56418958)), + ( + np.arange(5), + np.arange(1, 6), + np.arange(1, 6), + None, + (1.0, 3.12837917, 5.14094894, 7.02775903, 8.87030861), + ), + ( + np.arange(5), + np.arange(1, 6), + np.arange(1, 6), + (2, 5), + np.full((2, 5), (1.0, 3.12837917, 5.14094894, 7.02775903, 8.87030861)), + ), + ], +) +def test_skewnormal_moment(alpha, mu, sigma, size, expected): + with Model() as model: + SkewNormal("x", alpha=alpha, mu=mu, sigma=sigma, size=size) + assert_moment_is_expected(model, expected) + + +@pytest.mark.parametrize( + "b, kappa, mu, size, expected", + [ + (1.0, 1.0, 1.0, None, 1.0), + (1.0, np.ones(5), 1.0, None, np.full(5, 1.0)), + (np.arange(1, 6), 1.0, np.ones(5), None, np.full(5, 1.0)), + ( + np.arange(1, 6), + np.arange(1, 6), + np.arange(1, 6), + None, + (1.0, 1.25, 2.111111111111111, 3.0625, 4.04), + ), + ( + np.arange(1, 6), + np.arange(1, 6), + np.arange(1, 6), + (2, 5), + np.full((2, 5), (1.0, 1.25, 2.111111111111111, 3.0625, 4.04)), + ), + ], +) +def test_asymmetriclaplace_moment(b, kappa, mu, size, expected): + with Model() as model: + AsymmetricLaplace("x", b=b, kappa=kappa, mu=mu, size=size) + assert_moment_is_expected(model, expected)