Add back some missing moments (#5147)

* Exponential

* Laplace

* StudentT

* Cauchy

* Kumaraswamy

Author: michaeloriordan

pymc/distributions/continuous.py

@@ -1313,6 +1313,12 @@ def dist(cls, a, b, *args, **kwargs):
 
         return super().dist([a, b], *args, **kwargs)
 
+    def get_moment(rv, size, a, b):
+        mean = at.exp(at.log(b) + at.gammaln(1 + 1 / a) + at.gammaln(b) - at.gammaln(1 + 1 / a + b))
+        if not rv_size_is_none(size):
+            mean = at.full(size, mean)
+        return mean
+
     def logp(value, a, b):
         """
         Calculate log-probability of Kumaraswamy distribution at specified value.
@@ -1399,6 +1405,11 @@ def dist(cls, lam, *args, **kwargs):
         # Aesara exponential op is parametrized in terms of mu (1/lam)
         return super().dist([at.inv(lam)], **kwargs)
 
+    def get_moment(rv, size, mu):
+        if not rv_size_is_none(size):
+            mu = at.full(size, mu)
+        return mu
+
     def logcdf(value, mu):
         r"""
         Compute the log of cumulative distribution function for the Exponential distribution
@@ -1475,6 +1486,12 @@ def dist(cls, mu, b, *args, **kwargs):
         assert_negative_support(b, "b", "Laplace")
         return super().dist([mu, b], *args, **kwargs)
 
+    def get_moment(rv, size, mu, b):
+        mu, _ = at.broadcast_arrays(mu, b)
+        if not rv_size_is_none(size):
+            mu = at.full(size, mu)
+        return mu
+
     def logcdf(value, mu, b):
         """
         Compute the log of the cumulative distribution function for Laplace distribution
@@ -1800,6 +1817,12 @@ def dist(cls, nu, mu=0, lam=None, sigma=None, sd=None, *args, **kwargs):
 
         return super().dist([nu, mu, sigma], **kwargs)
 
+    def get_moment(rv, size, nu, mu, sigma):
+        mu, _, _ = at.broadcast_arrays(mu, nu, sigma)
+        if not rv_size_is_none(size):
+            mu = at.full(size, mu)
+        return mu
+
     def logp(value, nu, mu, sigma):
         """
         Calculate log-probability of StudentT distribution at specified value.
@@ -2001,12 +2024,15 @@ def dist(cls, alpha, beta, *args, **kwargs):
         alpha = at.as_tensor_variable(floatX(alpha))
         beta = at.as_tensor_variable(floatX(beta))
 
-        # median = alpha
-        # mode = alpha
-
         assert_negative_support(beta, "beta", "Cauchy")
         return super().dist([alpha, beta], **kwargs)
 
+    def get_moment(rv, size, alpha, beta):
+        alpha, _ = at.broadcast_arrays(alpha, beta)
+        if not rv_size_is_none(size):
+            alpha = at.full(size, alpha)
+        return alpha
+
     def logcdf(value, alpha, beta):
         """
         Compute the log of the cumulative distribution function for Cauchy distribution

pymc/tests/test_distributions_moments.py

@@ -2,7 +2,15 @@
 import pytest
 
 from pymc import Bernoulli, Flat, HalfFlat, Normal, TruncatedNormal, Uniform
-from pymc.distributions import Beta, HalfNormal
+from pymc.distributions import (
+    Beta,
+    Cauchy,
+    Exponential,
+    HalfNormal,
+    Kumaraswamy,
+    Laplace,
+    StudentT,
+)
 from pymc.distributions.shape_utils import rv_size_is_none
 from pymc.initial_point import make_initial_point_fn
 from pymc.model import Model
@@ -157,3 +165,79 @@ def test_beta_moment(alpha, beta, size, expected):
     with Model() as model:
         Beta("x", alpha=alpha, beta=beta, size=size)
     assert_moment_is_expected(model, expected)
+
+
+@pytest.mark.parametrize(
+    "lam, size, expected",
+    [
+        (2, None, 0.5),
+        (2, 5, np.full(5, 0.5)),
+        (np.arange(1, 5), None, 1 / np.arange(1, 5)),
+        (np.arange(1, 5), (2, 4), np.full((2, 4), 1 / np.arange(1, 5))),
+    ],
+)
+def test_exponential_moment(lam, size, expected):
+    with Model() as model:
+        Exponential("x", lam=lam, size=size)
+    assert_moment_is_expected(model, expected)
+
+
+@pytest.mark.parametrize(
+    "mu, b, size, expected",
+    [
+        (0, 1, None, 0),
+        (0, np.ones(5), None, np.zeros(5)),
+        (np.arange(5), 1, None, np.arange(5)),
+        (np.arange(5), np.arange(1, 6), (2, 5), np.full((2, 5), np.arange(5))),
+    ],
+)
+def test_laplace_moment(mu, b, size, expected):
+    with Model() as model:
+        Laplace("x", mu=mu, b=b, size=size)
+    assert_moment_is_expected(model, expected)
+
+
+@pytest.mark.parametrize(
+    "mu, nu, sigma, size, expected",
+    [
+        (0, 1, 1, None, 0),
+        (0, np.ones(5), 1, None, np.zeros(5)),
+        (np.arange(5), 10, np.arange(1, 6), None, np.arange(5)),
+        (np.arange(5), 10, np.arange(1, 6), (2, 5), np.full((2, 5), np.arange(5))),
+    ],
+)
+def test_studentt_moment(mu, nu, sigma, size, expected):
+    with Model() as model:
+        StudentT("x", mu=mu, nu=nu, sigma=sigma, size=size)
+    assert_moment_is_expected(model, expected)
+
+
+@pytest.mark.parametrize(
+    "alpha, beta, size, expected",
+    [
+        (0, 1, None, 0),
+        (0, np.ones(5), None, np.zeros(5)),
+        (np.arange(5), 1, None, np.arange(5)),
+        (np.arange(5), np.arange(1, 6), (2, 5), np.full((2, 5), np.arange(5))),
+    ],
+)
+def test_cauchy_moment(alpha, beta, size, expected):
+    with Model() as model:
+        Cauchy("x", alpha=alpha, beta=beta, size=size)
+    assert_moment_is_expected(model, expected)
+
+
+@pytest.mark.parametrize(
+    "a, b, size, expected",
+    [
+        (1, 1, None, 0.5),
+        (1, 1, 5, np.full(5, 0.5)),
+        (1, np.arange(1, 6), None, 1 / np.arange(2, 7)),
+        (np.arange(1, 6), 1, None, np.arange(1, 6) / np.arange(2, 7)),
+        (1, np.arange(1, 6), (2, 5), np.full((2, 5), 1 / np.arange(2, 7))),
+    ],
+)
+def test_kumaraswamy_moment(a, b, size, expected):
+    with Model() as model:
+        Kumaraswamy("x", a=a, b=b, size=size)
+    assert_moment_is_expected(model, expected)