NegativeBinomial, ZeroInflatedPoisson, ZeroInflatedBinomial moments (#5163)

* zero inflated binomial moment
* negative binomial moment
* ZIP moment

Author: theorashid

pymc/distributions/discrete.py

@@ -716,6 +716,12 @@ def get_n_p(cls, mu=None, alpha=None, p=None, n=None):
 
         return n, p
 
+    def get_moment(rv, size, n, p):
+        mu = at.floor(n * (1 - p) / p)
+        if not rv_size_is_none(size):
+            mu = at.full(size, mu)
+        return mu
+
     def logp(value, n, p):
         r"""
         Calculate log-probability of NegativeBinomial distribution at specified value.
@@ -1316,6 +1322,12 @@ def dist(cls, psi, theta, *args, **kwargs):
         theta = at.as_tensor_variable(floatX(theta))
         return super().dist([psi, theta], *args, **kwargs)
 
+    def get_moment(rv, size, psi, theta):
+        mean = at.floor(psi * theta)
+        if not rv_size_is_none(size):
+            mean = at.full(size, mean)
+        return mean
+
     def logp(value, psi, theta):
         r"""
         Calculate log-probability of ZeroInflatedPoisson distribution at specified value.
@@ -1449,6 +1461,12 @@ def dist(cls, psi, n, p, *args, **kwargs):
         p = at.as_tensor_variable(floatX(p))
         return super().dist([psi, n, p], *args, **kwargs)
 
+    def get_moment(rv, size, psi, n, p):
+        mean = at.round((1 - psi) * n * p)
+        if not rv_size_is_none(size):
+            mean = at.full(size, mean)
+        return mean
+
     def logp(value, psi, n, p):
         r"""
         Calculate log-probability of ZeroInflatedBinomial distribution at specified value.

pymc/tests/test_distributions_moments.py

@@ -19,9 +19,12 @@
     Laplace,
     Logistic,
     LogNormal,
+    NegativeBinomial,
     Poisson,
     StudentT,
     Weibull,
+    ZeroInflatedBinomial,
+    ZeroInflatedPoisson,
 )
 from pymc.distributions.shape_utils import rv_size_is_none
 from pymc.initial_point import make_initial_point_fn
@@ -402,6 +405,21 @@ def test_poisson_moment(mu, size, expected):
     assert_moment_is_expected(model, expected)
 
 
+@pytest.mark.parametrize(
+    "n, p, size, expected",
+    [
+        (10, 0.7, None, 4),
+        (10, 0.7, 5, np.full(5, 4)),
+        (np.full(3, 10), np.arange(1, 4) / 10, None, np.array([90, 40, 23])),
+        (10, np.arange(1, 4) / 10, (2, 3), np.full((2, 3), np.array([90, 40, 23]))),
+    ],
+)
+def test_negative_binomial_moment(n, p, size, expected):
+    with Model() as model:
+        NegativeBinomial("x", n=n, p=p, size=size)
+    assert_moment_is_expected(model, expected)
+
+
 @pytest.mark.parametrize(
     "c, size, expected",
     [
@@ -416,6 +434,36 @@ def test_constant_moment(c, size, expected):
     assert_moment_is_expected(model, expected)
 
 
+@pytest.mark.parametrize(
+    "psi, theta, size, expected",
+    [
+        (0.9, 3.0, None, 2),
+        (0.8, 2.9, 5, np.full(5, 2)),
+        (0.2, np.arange(1, 5) * 5, None, np.arange(1, 5)),
+        (0.2, np.arange(1, 5) * 5, (2, 4), np.full((2, 4), np.arange(1, 5))),
+    ],
+)
+def test_zero_inflated_poisson_moment(psi, theta, size, expected):
+    with Model() as model:
+        ZeroInflatedPoisson("x", psi=psi, theta=theta, size=size)
+    assert_moment_is_expected(model, expected)
+
+
+@pytest.mark.parametrize(
+    "psi, n, p, size, expected",
+    [
+        (0.2, 7, 0.7, None, 4),
+        (0.2, 7, 0.3, 5, np.full(5, 2)),
+        (0.6, 25, np.arange(1, 6) / 10, None, np.arange(1, 6)),
+        (0.6, 25, np.arange(1, 6) / 10, (2, 5), np.full((2, 5), np.arange(1, 6))),
+    ],
+)
+def test_zero_inflated_binomial_moment(psi, n, p, size, expected):
+    with Model() as model:
+        ZeroInflatedBinomial("x", psi=psi, n=n, p=p, size=size)
+    assert_moment_is_expected(model, expected)
+
+
 @pytest.mark.parametrize(
     "mu, s, size, expected",
     [