From ce508755ad9afc2e0a24cb658eb19694e81212de Mon Sep 17 00:00:00 2001 From: theorashid Date: Mon, 8 Nov 2021 23:36:33 +0000 Subject: [PATCH 1/5] zero inflated binomial moment --- pymc/distributions/discrete.py | 6 ++++++ pymc/tests/test_distributions_moments.py | 1 + 2 files changed, 7 insertions(+) diff --git a/pymc/distributions/discrete.py b/pymc/distributions/discrete.py index bb9ff075aa..4d6263e908 100644 --- a/pymc/distributions/discrete.py +++ b/pymc/distributions/discrete.py @@ -1449,6 +1449,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..99bfc5c50c 100644 --- a/pymc/tests/test_distributions_moments.py +++ b/pymc/tests/test_distributions_moments.py @@ -22,6 +22,7 @@ Poisson, StudentT, Weibull, + ZeroInflatedBinomial, ) from pymc.distributions.shape_utils import rv_size_is_none from pymc.initial_point import make_initial_point_fn From 537c1d305563eab4c11ebfa83e667b2ed651677d Mon Sep 17 00:00:00 2001 From: theorashid Date: Tue, 9 Nov 2021 00:07:31 +0000 Subject: [PATCH 2/5] negative binomial moment --- pymc/distributions/discrete.py | 6 ++++++ pymc/tests/test_distributions_moments.py | 20 ++++++++++++++++++++ 2 files changed, 26 insertions(+) diff --git a/pymc/distributions/discrete.py b/pymc/distributions/discrete.py index 4d6263e908..b8451ef1b6 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. diff --git a/pymc/tests/test_distributions_moments.py b/pymc/tests/test_distributions_moments.py index 99bfc5c50c..a97de9b9c1 100644 --- a/pymc/tests/test_distributions_moments.py +++ b/pymc/tests/test_distributions_moments.py @@ -19,6 +19,7 @@ Laplace, Logistic, LogNormal, + NegativeBinomial, Poisson, StudentT, Weibull, @@ -404,7 +405,26 @@ def test_poisson_moment(mu, size, expected): @pytest.mark.parametrize( +<<<<<<< HEAD "c, size, expected", +======= + "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( + "psi, n, p, size, expected", +>>>>>>> 2eb37150 (negative binomial moment) [ (1, None, 1), (1, 5, np.full(5, 1)), From f0fe7a1ca108cde4456fa412e9dbe511deafd92f Mon Sep 17 00:00:00 2001 From: theorashid Date: Tue, 9 Nov 2021 09:59:05 +0000 Subject: [PATCH 3/5] ZIP moment --- pymc/distributions/discrete.py | 6 ++++++ pymc/tests/test_distributions_moments.py | 16 ++++++++++++++++ 2 files changed, 22 insertions(+) diff --git a/pymc/distributions/discrete.py b/pymc/distributions/discrete.py index b8451ef1b6..de2b11b606 100644 --- a/pymc/distributions/discrete.py +++ b/pymc/distributions/discrete.py @@ -1322,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): + mu = at.floor(psi * theta) + if not rv_size_is_none(size): + mu = at.full(size, mu) + return mu + def logp(value, psi, theta): r""" Calculate log-probability of ZeroInflatedPoisson distribution at specified value. diff --git a/pymc/tests/test_distributions_moments.py b/pymc/tests/test_distributions_moments.py index a97de9b9c1..94db912251 100644 --- a/pymc/tests/test_distributions_moments.py +++ b/pymc/tests/test_distributions_moments.py @@ -24,6 +24,7 @@ StudentT, Weibull, ZeroInflatedBinomial, + ZeroInflatedPoisson, ) from pymc.distributions.shape_utils import rv_size_is_none from pymc.initial_point import make_initial_point_fn @@ -422,6 +423,21 @@ def test_negative_binomial_moment(n, p, 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", >>>>>>> 2eb37150 (negative binomial moment) From 5848d86704e380be488ffe2ab913d0f4e27f75ba Mon Sep 17 00:00:00 2001 From: theorashid Date: Tue, 9 Nov 2021 12:04:31 +0000 Subject: [PATCH 4/5] fix after rebase --- pymc/tests/test_distributions_moments.py | 31 ++++++++++++++++-------- 1 file changed, 21 insertions(+), 10 deletions(-) diff --git a/pymc/tests/test_distributions_moments.py b/pymc/tests/test_distributions_moments.py index 94db912251..ca5337c372 100644 --- a/pymc/tests/test_distributions_moments.py +++ b/pymc/tests/test_distributions_moments.py @@ -406,9 +406,6 @@ def test_poisson_moment(mu, size, expected): @pytest.mark.parametrize( -<<<<<<< HEAD - "c, size, expected", -======= "n, p, size, expected", [ (10, 0.7, None, 4), @@ -423,6 +420,20 @@ def test_negative_binomial_moment(n, p, size, expected): assert_moment_is_expected(model, expected) +@pytest.mark.parametrize( + "c, size, expected", + [ + (1, None, 1), + (1, 5, np.full(5, 1)), + (np.arange(1, 6), None, np.arange(1, 6)), + ], +) +def test_constant_moment(c, size, expected): + with Model() as model: + Constant("x", c=c, size=size) + assert_moment_is_expected(model, expected) + + @pytest.mark.parametrize( "psi, theta, size, expected", [ @@ -435,21 +446,21 @@ def test_negative_binomial_moment(n, p, size, expected): 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) + assert_moment_is_expected(model, expected) @pytest.mark.parametrize( "psi, n, p, size, expected", ->>>>>>> 2eb37150 (negative binomial moment) [ - (1, None, 1), - (1, 5, np.full(5, 1)), - (np.arange(1, 6), None, np.arange(1, 6)), + (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_constant_moment(c, size, expected): +def test_zero_inflated_binomial_moment(psi, n, p, size, expected): with Model() as model: - Constant("x", c=c, size=size) + ZeroInflatedBinomial("x", psi=psi, n=n, p=p, size=size) assert_moment_is_expected(model, expected) From 1381155b1d3bef25b7cb38aa3eddd456b29e30c9 Mon Sep 17 00:00:00 2001 From: theorashid Date: Tue, 9 Nov 2021 14:49:35 +0000 Subject: [PATCH 5/5] change mu to mean in ZIP --- pymc/distributions/discrete.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pymc/distributions/discrete.py b/pymc/distributions/discrete.py index de2b11b606..31f2dd7d49 100644 --- a/pymc/distributions/discrete.py +++ b/pymc/distributions/discrete.py @@ -1323,10 +1323,10 @@ def dist(cls, psi, theta, *args, **kwargs): return super().dist([psi, theta], *args, **kwargs) def get_moment(rv, size, psi, theta): - mu = at.floor(psi * theta) + mean = at.floor(psi * theta) if not rv_size_is_none(size): - mu = at.full(size, mu) - return mu + mean = at.full(size, mean) + return mean def logp(value, psi, theta): r"""