From cd9f9d12d8353b3cead18f9fba45bf7bba6bcddf Mon Sep 17 00:00:00 2001 From: Farhan Reynaldo Date: Sun, 7 Nov 2021 16:10:09 +0700 Subject: [PATCH 1/4] add poisson and binomial moment --- pymc/distributions/discrete.py | 11 ++++++ pymc/tests/test_distributions_moments.py | 45 ++++++++++++++++++++++-- 2 files changed, 54 insertions(+), 2 deletions(-) diff --git a/pymc/distributions/discrete.py b/pymc/distributions/discrete.py index 35fd2bc2f6..51d5deb595 100644 --- a/pymc/distributions/discrete.py +++ b/pymc/distributions/discrete.py @@ -117,6 +117,12 @@ def dist(cls, n, p, *args, **kwargs): # mode = at.cast(tround(n * p), self.dtype) return super().dist([n, p], **kwargs) + def get_moment(rv, size, n, p): + mean = n * p + if not rv_size_is_none(size): + mean = at.full(size, mean) + return mean + def logp(value, n, p): r""" Calculate log-probability of Binomial distribution at specified value. @@ -570,6 +576,11 @@ def dist(cls, mu, *args, **kwargs): # mode = intX(at.floor(mu)) return super().dist([mu], *args, **kwargs) + def get_moment(rv, size, mu): + if not rv_size_is_none(size): + mu = at.full(size, mu) + return mu + def logp(value, mu): r""" Calculate log-probability of Poisson distribution at specified value. diff --git a/pymc/tests/test_distributions_moments.py b/pymc/tests/test_distributions_moments.py index 36fe591e69..5ff7333f69 100644 --- a/pymc/tests/test_distributions_moments.py +++ b/pymc/tests/test_distributions_moments.py @@ -6,6 +6,7 @@ from pymc import Bernoulli, Flat, HalfFlat, Normal, TruncatedNormal, Uniform from pymc.distributions import ( Beta, + Binomial, Cauchy, Exponential, Gamma, @@ -14,6 +15,7 @@ Kumaraswamy, Laplace, LogNormal, + Poisson, StudentT, Weibull, ) @@ -209,7 +211,13 @@ def test_laplace_moment(mu, b, 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))), + ( + 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): @@ -318,7 +326,10 @@ def test_gamma_moment(alpha, beta, size, expected): np.arange(1, 6), np.arange(2, 7), (2, 5), - np.full((2, 5), np.arange(2, 7) * special.gamma(1 + 1 / np.arange(1, 6))), + np.full( + (2, 5), + np.arange(2, 7) * special.gamma(1 + 1 / np.arange(1, 6)), + ), ), ], ) @@ -326,3 +337,33 @@ def test_weibull_moment(alpha, beta, size, expected): with Model() as model: Weibull("x", alpha=alpha, beta=beta, size=size) assert_moment_is_expected(model, expected) + + +@pytest.mark.parametrize( + "n, p, size, expected", + [ + (10, 0.5, None, 5), + (10, 0.5, 5, np.full(5, 5)), + (10, np.arange(1, 6) / 10, None, np.arange(1, 6)), + (10, np.arange(1, 6) / 10, (2, 5), np.full((2, 5), np.arange(1, 6))), + ], +) +def test_binomial_moment(n, p, size, expected): + with Model() as model: + Binomial("x", n=n, p=p, size=size) + assert_moment_is_expected(model, expected) + + +@pytest.mark.parametrize( + "mu, size, expected", + [ + (2, None, 2), + (2, 5, np.full(5, 2)), + (np.arange(1, 5), None, np.arange(1, 5)), + (np.arange(1, 5), (2, 4), np.full((2, 4), np.arange(1, 5))), + ], +) +def test_poisson_moment(mu, size, expected): + with Model() as model: + Poisson("x", mu=mu, size=size) + assert_moment_is_expected(model, expected) From 81d682ec72a4291f417b6b6b660a40ee28d3aa6c Mon Sep 17 00:00:00 2001 From: Farhan Reynaldo Date: Mon, 8 Nov 2021 21:05:38 +0700 Subject: [PATCH 2/4] remove commented code, floor and round moment --- pymc/distributions/discrete.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/pymc/distributions/discrete.py b/pymc/distributions/discrete.py index 51d5deb595..d8a8c2f8fd 100644 --- a/pymc/distributions/discrete.py +++ b/pymc/distributions/discrete.py @@ -114,11 +114,10 @@ class Binomial(Discrete): def dist(cls, n, p, *args, **kwargs): n = at.as_tensor_variable(intX(n)) p = at.as_tensor_variable(floatX(p)) - # mode = at.cast(tround(n * p), self.dtype) return super().dist([n, p], **kwargs) def get_moment(rv, size, n, p): - mean = n * p + mean = at.round(n * p) if not rv_size_is_none(size): mean = at.full(size, mean) return mean @@ -573,10 +572,10 @@ class Poisson(Discrete): @classmethod def dist(cls, mu, *args, **kwargs): mu = at.as_tensor_variable(floatX(mu)) - # mode = intX(at.floor(mu)) return super().dist([mu], *args, **kwargs) def get_moment(rv, size, mu): + mu = at.floor(mu) if not rv_size_is_none(size): mu = at.full(size, mu) return mu From 93fe90153e030bda44357516cb3d0ad343cc3b97 Mon Sep 17 00:00:00 2001 From: Farhan Reynaldo Date: Mon, 8 Nov 2021 21:26:46 +0700 Subject: [PATCH 3/4] non integer values on mu for robust test --- pymc/tests/test_distributions_moments.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pymc/tests/test_distributions_moments.py b/pymc/tests/test_distributions_moments.py index 5ff7333f69..74483226e3 100644 --- a/pymc/tests/test_distributions_moments.py +++ b/pymc/tests/test_distributions_moments.py @@ -357,8 +357,8 @@ def test_binomial_moment(n, p, size, expected): @pytest.mark.parametrize( "mu, size, expected", [ - (2, None, 2), - (2, 5, np.full(5, 2)), + (2.7, None, 2), + (2.3, 5, np.full(5, 2)), (np.arange(1, 5), None, np.arange(1, 5)), (np.arange(1, 5), (2, 4), np.full((2, 4), np.arange(1, 5))), ], From 57789b17cf9cfb6ef5daf466b9fc60d288b576f6 Mon Sep 17 00:00:00 2001 From: Farhan Reynaldo Date: Mon, 8 Nov 2021 21:31:06 +0700 Subject: [PATCH 4/4] different rounded binomial for robust test --- pymc/tests/test_distributions_moments.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pymc/tests/test_distributions_moments.py b/pymc/tests/test_distributions_moments.py index 74483226e3..f6eb8c5f1b 100644 --- a/pymc/tests/test_distributions_moments.py +++ b/pymc/tests/test_distributions_moments.py @@ -342,8 +342,8 @@ def test_weibull_moment(alpha, beta, size, expected): @pytest.mark.parametrize( "n, p, size, expected", [ - (10, 0.5, None, 5), - (10, 0.5, 5, np.full(5, 5)), + (7, 0.7, None, 5), + (7, 0.3, 5, np.full(5, 2)), (10, np.arange(1, 6) / 10, None, np.arange(1, 6)), (10, np.arange(1, 6) / 10, (2, 5), np.full((2, 5), np.arange(1, 6))), ],