From d500aecde19b27453eed4b8186dc2a672b64890b Mon Sep 17 00:00:00 2001 From: Morgan Pihl Date: Thu, 11 Nov 2021 21:09:13 +0100 Subject: [PATCH 1/3] adding beta-binomial mean and test cases. Note that the get_moment method returns an integer value (rounding the mean down to the nearest integer) in this implementation --- 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 86e7f1e9b9..df9615ef4b 100644 --- a/pymc/distributions/discrete.py +++ b/pymc/distributions/discrete.py @@ -238,6 +238,12 @@ def dist(cls, alpha, beta, n, *args, **kwargs): beta = at.as_tensor_variable(floatX(beta)) n = at.as_tensor_variable(intX(n)) return super().dist([n, alpha, beta], **kwargs) + + def get_moment(rv, size, n, alpha, beta): + mean = (n * alpha) / (alpha + beta) + if not rv_size_is_none(size): + mean = at.full(size, mean) + return mean def logp(value, n, alpha, beta): r""" diff --git a/pymc/tests/test_distributions_moments.py b/pymc/tests/test_distributions_moments.py index bbe478f8ac..4adbbd4735 100644 --- a/pymc/tests/test_distributions_moments.py +++ b/pymc/tests/test_distributions_moments.py @@ -6,6 +6,7 @@ from pymc.distributions import ( Bernoulli, Beta, + BetaBinomial, Binomial, Cauchy, ChiSquared, @@ -209,6 +210,21 @@ def test_beta_moment(alpha, beta, size, expected): assert_moment_is_expected(model, expected) +@pytest.mark.parametrize( + "n, alpha, beta, size, expected", + [ + (10, 1, 1, None, 5), + (10, 1, 1, 5, np.full(5, 5)), + (10, 1, np.arange(1, 6), None, np.floor(10 / np.arange(2, 7))), + (10, 1, np.arange(1, 6), (2, 5), np.full((2, 5), np.floor(10 / np.arange(2, 7)))), + ], +) +def test_beta_binomial_moment(alpha, beta, n, size, expected): + with Model() as model: + BetaBinomial("x", alpha=alpha, beta=beta, n=n, size=size) + assert_moment_is_expected(model, expected) + + @pytest.mark.parametrize( "nu, size, expected", [ From 8c29645778bfcb7f055720feb60ddf18f9fb8196 Mon Sep 17 00:00:00 2001 From: Morgan Pihl Date: Fri, 12 Nov 2021 11:37:30 +0100 Subject: [PATCH 2/3] removes trailing whitespace --- pymc/distributions/discrete.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pymc/distributions/discrete.py b/pymc/distributions/discrete.py index df9615ef4b..4b1b01b3f0 100644 --- a/pymc/distributions/discrete.py +++ b/pymc/distributions/discrete.py @@ -238,7 +238,7 @@ def dist(cls, alpha, beta, n, *args, **kwargs): beta = at.as_tensor_variable(floatX(beta)) n = at.as_tensor_variable(intX(n)) return super().dist([n, alpha, beta], **kwargs) - + def get_moment(rv, size, n, alpha, beta): mean = (n * alpha) / (alpha + beta) if not rv_size_is_none(size): From 2d8bf027b221ebaef1bf3483d61571222efcad8f Mon Sep 17 00:00:00 2001 From: Morgan Pihl Date: Sat, 13 Nov 2021 09:15:16 +0100 Subject: [PATCH 3/3] adds explicit rounding to get_moment --- pymc/distributions/discrete.py | 2 +- pymc/tests/test_distributions_moments.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pymc/distributions/discrete.py b/pymc/distributions/discrete.py index 4b1b01b3f0..e190bf7a12 100644 --- a/pymc/distributions/discrete.py +++ b/pymc/distributions/discrete.py @@ -240,7 +240,7 @@ def dist(cls, alpha, beta, n, *args, **kwargs): return super().dist([n, alpha, beta], **kwargs) def get_moment(rv, size, n, alpha, beta): - mean = (n * alpha) / (alpha + beta) + mean = at.round((n * alpha) / (alpha + beta)) if not rv_size_is_none(size): mean = at.full(size, mean) return mean diff --git a/pymc/tests/test_distributions_moments.py b/pymc/tests/test_distributions_moments.py index 4adbbd4735..20ea2bc2f4 100644 --- a/pymc/tests/test_distributions_moments.py +++ b/pymc/tests/test_distributions_moments.py @@ -215,8 +215,8 @@ def test_beta_moment(alpha, beta, size, expected): [ (10, 1, 1, None, 5), (10, 1, 1, 5, np.full(5, 5)), - (10, 1, np.arange(1, 6), None, np.floor(10 / np.arange(2, 7))), - (10, 1, np.arange(1, 6), (2, 5), np.full((2, 5), np.floor(10 / np.arange(2, 7)))), + (10, 1, np.arange(1, 6), None, np.round(10 / np.arange(2, 7))), + (10, 1, np.arange(1, 6), (2, 5), np.full((2, 5), np.round(10 / np.arange(2, 7)))), ], ) def test_beta_binomial_moment(alpha, beta, n, size, expected):