Adding moments for AsymmetricLaplace and SkewNormal and corresponding tests

Author: juan.lopez.arriaza

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.

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)