add hyper geometric moment

Author: farhanreynaldo

pymc/distributions/discrete.py

@@ -926,6 +926,12 @@ def dist(cls, N, k, n, *args, **kwargs):
         n = at.as_tensor_variable(intX(n))
         return super().dist([good, bad, n], *args, **kwargs)
 
+    def get_moment(rv, size, N, k, n):
+        mode = intX(at.floor((n + 1) * (k + 1) / (N + 2)))
+        if not rv_size_is_none(size):
+            mode = at.full(size, mode)
+        return mode
+
     def logp(value, good, bad, n):
         r"""
         Calculate log-probability of HyperGeometric distribution at specified value.

pymc/tests/test_distributions_moments.py

@@ -18,6 +18,7 @@
     HalfFlat,
     HalfNormal,
     HalfStudentT,
+    HyperGeometric,
     Kumaraswamy,
     Laplace,
     Logistic,
@@ -157,7 +158,9 @@ def test_halfstudentt_moment(nu, sigma, size, expected):
     assert_moment_is_expected(model, expected)
 
 
-@pytest.mark.skip(reason="aeppl interval transform fails when both edges are None")
+@pytest.mark.skip(
+    reason="aeppl interval transform fails when both edges are None"
+)
 @pytest.mark.parametrize(
     "mu, sigma, lower, upper, size, expected",
     [
@@ -169,7 +172,9 @@ def test_halfstudentt_moment(nu, sigma, size, expected):
 )
 def test_truncatednormal_moment(mu, sigma, lower, upper, size, expected):
     with Model() as model:
-        TruncatedNormal("x", mu=mu, sigma=sigma, lower=lower, upper=upper, size=size)
+        TruncatedNormal(
+            "x", mu=mu, sigma=sigma, lower=lower, upper=upper, size=size
+        )
     assert_moment_is_expected(model, expected)
 
 
@@ -417,7 +422,12 @@ def test_poisson_moment(mu, 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]))),
+        (
+            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):
@@ -461,7 +471,13 @@ def test_zero_inflated_poisson_moment(psi, theta, 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))),
+        (
+            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):
@@ -503,3 +519,24 @@ def test_geometric_moment(p, size, expected):
     with Model() as model:
         Geometric("x", p=p, size=size)
     assert_moment_is_expected(model, expected)
+
+
+@pytest.mark.parametrize(
+    "N, k, n, size, expected",
+    [
+        (50, 10, 20, None, 4),
+        (50, 10, 23, 5, np.full(5, 5)),
+        (50, 10, np.arange(23, 28), None, np.full(5, 5)),
+        (
+            50,
+            10,
+            np.arange(18, 23),
+            (2, 5),
+            np.full((2, 5), 4),
+        ),
+    ],
+)
+def test_hyper_geometric_moment(N, k, n, size, expected):
+    with Model() as model:
+        HyperGeometric("x", N=N, k=k, n=n, size=size)
+    assert_moment_is_expected(model, expected)