diff --git a/pymc/distributions/discrete.py b/pymc/distributions/discrete.py index bb9ff075aa..31f2dd7d49 100644 --- a/pymc/distributions/discrete.py +++ b/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. diff --git a/pymc/tests/test_distributions_moments.py b/pymc/tests/test_distributions_moments.py index 9e2d56afdb..ca5337c372 100644 --- a/pymc/tests/test_distributions_moments.py +++ b/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", [