Added moments for gumbel, triangular and logitnormal distributions for issue #5078 (#5180)

Author: sagartomar

pymc/distributions/continuous.py

@@ -86,7 +86,7 @@ def polyagamma_cdf(*args, **kwargs):
 )
 from pymc.distributions.distribution import Continuous
 from pymc.distributions.shape_utils import rv_size_is_none
-from pymc.math import logdiffexp, logit
+from pymc.math import invlogit, logdiffexp, logit
 from pymc.util import UNSET
 
 __all__ = [
@@ -3101,6 +3101,12 @@ def dist(cls, lower=0, upper=1, c=0.5, *args, **kwargs):
 
         return super().dist([lower, c, upper], *args, **kwargs)
 
+    def get_moment(rv, size, lower, c, upper):
+        mean = (lower + upper + c) / 3
+        if not rv_size_is_none(size):
+            mean = at.full(size, mean)
+        return mean
+
     def logcdf(value, lower, c, upper):
         """
         Compute the log of the cumulative distribution function for Triangular distribution
@@ -3198,6 +3204,12 @@ def dist(
 
         return super().dist([mu, beta], **kwargs)
 
+    def get_moment(rv, size, mu, beta):
+        mean = mu + beta * np.euler_gamma
+        if not rv_size_is_none(size):
+            mean = at.full(size, mean)
+        return mean
+
     def _distr_parameters_for_repr(self):
         return ["mu", "beta"]
 
@@ -3501,6 +3513,12 @@ def dist(cls, mu=0, sigma=None, tau=None, sd=None, **kwargs):
 
         return super().dist([mu, sigma], **kwargs)
 
+    def get_moment(rv, size, mu, sigma):
+        median, _ = at.broadcast_arrays(invlogit(mu), sigma)
+        if not rv_size_is_none(size):
+            median = at.full(size, median)
+        return median
+
     def logp(value, mu, sigma):
         """
         Calculate log-probability of LogitNormal distribution at specified value.

pymc/tests/test_distributions_moments.py

@@ -17,6 +17,7 @@
     Flat,
     Gamma,
     Geometric,
+    Gumbel,
     HalfCauchy,
     HalfFlat,
     HalfNormal,
@@ -25,12 +26,14 @@
     Kumaraswamy,
     Laplace,
     Logistic,
+    LogitNormal,
     LogNormal,
     NegativeBinomial,
     Normal,
     Pareto,
     Poisson,
     StudentT,
+    Triangular,
     TruncatedNormal,
     Uniform,
     Wald,
@@ -650,3 +653,62 @@ def test_dirichlet_moment(a, size, expected):
     with Model() as model:
         Dirichlet("x", a=a, size=size)
     assert_moment_is_expected(model, expected)
+
+
+@pytest.mark.parametrize(
+    "mu, beta, size, expected",
+    [
+        (0, 2, None, 2 * np.euler_gamma),
+        (1, np.arange(1, 4), None, 1 + np.arange(1, 4) * np.euler_gamma),
+        (np.arange(5), 2, None, np.arange(5) + 2 * np.euler_gamma),
+        (1, 2, 5, np.full(5, 1 + 2 * np.euler_gamma)),
+        (
+            np.arange(5),
+            np.arange(1, 6),
+            (2, 5),
+            np.full((2, 5), np.arange(5) + np.arange(1, 6) * np.euler_gamma),
+        ),
+    ],
+)
+def test_gumbel_moment(mu, beta, size, expected):
+    with Model() as model:
+        Gumbel("x", mu=mu, beta=beta, size=size)
+    assert_moment_is_expected(model, expected)
+
+
+@pytest.mark.parametrize(
+    "c, lower, upper, size, expected",
+    [
+        (1, 0, 5, None, 2),
+        (3, np.arange(-3, 6, 3), np.arange(3, 12, 3), None, np.array([1, 3, 5])),
+        (np.arange(-3, 6, 3), -3, 3, None, np.array([-1, 0, 1])),
+        (3, -3, 6, 5, np.full(5, 2)),
+        (
+            np.arange(-3, 6, 3),
+            np.arange(-9, -2, 3),
+            np.arange(3, 10, 3),
+            (2, 3),
+            np.full((2, 3), np.array([-3, 0, 3])),
+        ),
+    ],
+)
+def test_triangular_moment(c, lower, upper, size, expected):
+    with Model() as model:
+        Triangular("x", c=c, lower=lower, upper=upper, size=size)
+    assert_moment_is_expected(model, expected)
+
+
+@pytest.mark.parametrize(
+    "mu, sigma, size, expected",
+    [
+        (1, 2, None, special.expit(1)),
+        (0, np.arange(1, 5), None, special.expit(np.zeros(4))),
+        (np.arange(4), 1, None, special.expit(np.arange(4))),
+        (1, 5, 4, special.expit(np.ones(4))),
+        (np.arange(4), np.arange(1, 5), (2, 4), np.full((2, 4), special.expit(np.arange(4)))),
+    ],
+)
+def test_logitnormal_moment(mu, sigma, size, expected):
+    with Model() as model:
+        LogitNormal("x", mu=mu, sigma=sigma, size=size)
+    assert_moment_is_expected(model, expected)