From 549185f0ca697ad84f470a2c439dee79c6bc6ed4 Mon Sep 17 00:00:00 2001 From: Morgan Pihl Date: Thu, 18 Nov 2021 09:48:23 +0100 Subject: [PATCH 1/5] fixes broadcasting when size argument is present --- pymc/distributions/multivariate.py | 7 +++++++ pymc/tests/test_distributions_moments.py | 17 +++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/pymc/distributions/multivariate.py b/pymc/distributions/multivariate.py index 729d393c08..c33b5566a2 100644 --- a/pymc/distributions/multivariate.py +++ b/pymc/distributions/multivariate.py @@ -518,6 +518,13 @@ def dist(cls, n, p, *args, **kwargs): return super().dist([n, p], *args, **kwargs) + def get_moment(rv, size, n, p): + mode = at.round(n * p) + if not rv_size_is_none(size): + output_size = at.concatenate([size, p.shape]) + mode = at.full(output_size, mode) + return mode + def logp(value, n, p): """ Calculate log-probability of Multinomial distribution diff --git a/pymc/tests/test_distributions_moments.py b/pymc/tests/test_distributions_moments.py index 777be4a8bc..ef5a8a8915 100644 --- a/pymc/tests/test_distributions_moments.py +++ b/pymc/tests/test_distributions_moments.py @@ -37,6 +37,7 @@ LogitNormal, LogNormal, Moyal, + Multinomial, NegativeBinomial, Normal, Pareto, @@ -1036,3 +1037,19 @@ def test_polyagamma_moment(h, z, size, expected): with Model() as model: PolyaGamma("x", h=h, z=z, size=size) assert_moment_is_expected(model, expected) + + +@pytest.mark.parametrize( + "p, n, size, expected", + [ + (np.array([0.25, 0.25, 0.25, 0.25]), 1, None, np.repeat(0, 4)), + (np.array([0.3, 0.6, 0.05, 0.05]), 2, None, np.array([1, 1, 0, 0])), + (np.array([0.3, 0.6, 0.05, 0.05]), 10, None, np.array([3, 6, 0, 0])), + (np.array([[0.5, 0.5], [0.6, 0.4]]), np.array([1, 10]), None, np.array([[0, 5], [1, 4]])), + (np.array([[0.5, 0.5], [0.6, 0.4]]), np.array([1, 10]), 2, np.full((2, 2, 2), [[0, 5], [1, 4]])), + ], +) +def test_multinomial_moment(p, n, size, expected): + with Model() as model: + Multinomial("x", n=n, p=p, size=size) + assert_moment_is_expected(model, expected) From 504f853062a1eafa7416f46ebb0597f4df0b5a0c Mon Sep 17 00:00:00 2001 From: Morgan Pihl Date: Thu, 18 Nov 2021 09:54:14 +0100 Subject: [PATCH 2/5] removes old tests --- pymc/tests/test_distributions.py | 27 --------------------------- 1 file changed, 27 deletions(-) diff --git a/pymc/tests/test_distributions.py b/pymc/tests/test_distributions.py index 68822b264f..773fd85d75 100644 --- a/pymc/tests/test_distributions.py +++ b/pymc/tests/test_distributions.py @@ -2142,25 +2142,6 @@ def test_multinomial(self, n): Multinomial, Vector(Nat, n), {"p": Simplex(n), "n": Nat}, multinomial_logpdf ) - @pytest.mark.skip(reason="Moment calculations have not been refactored yet") - @pytest.mark.parametrize( - "p,n", - [ - [[0.25, 0.25, 0.25, 0.25], 1], - [[0.3, 0.6, 0.05, 0.05], 2], - [[0.3, 0.6, 0.05, 0.05], 10], - ], - ) - def test_multinomial_mode(self, p, n): - _p = np.array(p) - with Model() as model: - m = Multinomial("m", n, _p, _p.shape) - assert_allclose(m.distribution.mode.eval().sum(), n) - _p = np.array([p, p]) - with Model() as model: - m = Multinomial("m", n, _p, _p.shape) - assert_allclose(m.distribution.mode.eval().sum(axis=-1), n) - @pytest.mark.parametrize( "p, size, n", [ @@ -2188,14 +2169,6 @@ def test_multinomial_random(self, p, size, n): assert m.eval().shape == size + p.shape - @pytest.mark.skip(reason="Moment calculations have not been refactored yet") - def test_multinomial_mode_with_shape(self): - n = [1, 10] - p = np.asarray([[0.25, 0.25, 0.25, 0.25], [0.26, 0.26, 0.26, 0.22]]) - with Model() as model: - m = Multinomial("m", n=n, p=p, size=(2, 4)) - assert_allclose(m.distribution.mode.eval().sum(axis=-1), n) - def test_multinomial_vec(self): vals = np.array([[2, 4, 4], [3, 3, 4]]) p = np.array([0.2, 0.3, 0.5]) From 17ef35178db149e668c500dc572b7e9fd3a98ebf Mon Sep 17 00:00:00 2001 From: Morgan Pihl Date: Thu, 18 Nov 2021 11:40:56 +0100 Subject: [PATCH 3/5] changed test cases and treatment of dimensions. Last dimension represents outcomes/events, second from last dimension represents batches (independent multinomial distributions) and first dimension is set by the size argument. Note that this is different from previous effects of the size argument --- pymc/distributions/multivariate.py | 2 +- pymc/tests/test_distributions_moments.py | 13 +++++++++++-- 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/pymc/distributions/multivariate.py b/pymc/distributions/multivariate.py index c33b5566a2..1a894b7f08 100644 --- a/pymc/distributions/multivariate.py +++ b/pymc/distributions/multivariate.py @@ -519,7 +519,7 @@ def dist(cls, n, p, *args, **kwargs): return super().dist([n, p], *args, **kwargs) def get_moment(rv, size, n, p): - mode = at.round(n * p) + mode = at.round(n * p.transpose()).transpose() if not rv_size_is_none(size): output_size = at.concatenate([size, p.shape]) mode = at.full(output_size, mode) diff --git a/pymc/tests/test_distributions_moments.py b/pymc/tests/test_distributions_moments.py index ef5a8a8915..b34a91143a 100644 --- a/pymc/tests/test_distributions_moments.py +++ b/pymc/tests/test_distributions_moments.py @@ -1045,8 +1045,17 @@ def test_polyagamma_moment(h, z, size, expected): (np.array([0.25, 0.25, 0.25, 0.25]), 1, None, np.repeat(0, 4)), (np.array([0.3, 0.6, 0.05, 0.05]), 2, None, np.array([1, 1, 0, 0])), (np.array([0.3, 0.6, 0.05, 0.05]), 10, None, np.array([3, 6, 0, 0])), - (np.array([[0.5, 0.5], [0.6, 0.4]]), np.array([1, 10]), None, np.array([[0, 5], [1, 4]])), - (np.array([[0.5, 0.5], [0.6, 0.4]]), np.array([1, 10]), 2, np.full((2, 2, 2), [[0, 5], [1, 4]])), + ( + np.array([[0.25, 0.25, 0.25, 0.25], [0.26, 0.26, 0.26, 0.22]]), + np.array([1, 10]), + None, + np.array([[0, 0, 0, 0], [3, 3, 3, 2]]) + ), + ( + np.array([[0.25, 0.25, 0.25, 0.25], [0.26, 0.26, 0.26, 0.22]]), + np.array([1, 10]), + 2, + np.full((2, 2, 4), [[0, 0, 0, 0], [3, 3, 3, 2]])), ], ) def test_multinomial_moment(p, n, size, expected): From 49cdfe898ea5b8f6b949a99e655cd59fefedcebf Mon Sep 17 00:00:00 2001 From: Morgan Pihl Date: Wed, 24 Nov 2021 21:09:58 +0100 Subject: [PATCH 4/5] adds support for broadcasting dimensions --- pymc/distributions/multivariate.py | 10 +++++++++- pymc/tests/test_distributions_moments.py | 15 +++++++++++---- 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/pymc/distributions/multivariate.py b/pymc/distributions/multivariate.py index 1a894b7f08..cdb43d5684 100644 --- a/pymc/distributions/multivariate.py +++ b/pymc/distributions/multivariate.py @@ -519,7 +519,15 @@ def dist(cls, n, p, *args, **kwargs): return super().dist([n, p], *args, **kwargs) def get_moment(rv, size, n, p): - mode = at.round(n * p.transpose()).transpose() + if p.ndim > 1: + n = at.shape_padright(n) + if (p.ndim == 1) & (n.ndim > 0): + n = at.shape_padright(n) + p = at.shape_padleft(p) + mode = at.round(n * p) + diff = n - at.sum(mode, axis=-1, keepdims=True) + inc_bool_arr = at.abs_(diff) > 0 + mode = at.inc_subtensor(mode[inc_bool_arr.nonzero()], diff[inc_bool_arr.nonzero()]) if not rv_size_is_none(size): output_size = at.concatenate([size, p.shape]) mode = at.full(output_size, mode) diff --git a/pymc/tests/test_distributions_moments.py b/pymc/tests/test_distributions_moments.py index b34a91143a..92e12d0480 100644 --- a/pymc/tests/test_distributions_moments.py +++ b/pymc/tests/test_distributions_moments.py @@ -1042,20 +1042,27 @@ def test_polyagamma_moment(h, z, size, expected): @pytest.mark.parametrize( "p, n, size, expected", [ - (np.array([0.25, 0.25, 0.25, 0.25]), 1, None, np.repeat(0, 4)), + (np.array([0.25, 0.25, 0.25, 0.25]), 1, None, np.array([1, 0, 0, 0])), (np.array([0.3, 0.6, 0.05, 0.05]), 2, None, np.array([1, 1, 0, 0])), - (np.array([0.3, 0.6, 0.05, 0.05]), 10, None, np.array([3, 6, 0, 0])), + (np.array([0.3, 0.6, 0.05, 0.05]), 10, None, np.array([4, 6, 0, 0])), + (np.array([[0.3, 0.6, 0.05, 0.05], [0.25, 0.25, 0.25, 0.25]]), 10, None, np.array([[4, 6, 0, 0], [4, 2, 2, 2]])), ( np.array([[0.25, 0.25, 0.25, 0.25], [0.26, 0.26, 0.26, 0.22]]), np.array([1, 10]), None, - np.array([[0, 0, 0, 0], [3, 3, 3, 2]]) + np.array([[1, 0, 0, 0], [2, 3, 3, 2]]) + ), + ( + np.array([0.26, 0.26, 0.26, 0.22]), + np.array([1, 10]), + None, + np.array([[1, 0, 0, 0], [2, 3, 3, 2]]) ), ( np.array([[0.25, 0.25, 0.25, 0.25], [0.26, 0.26, 0.26, 0.22]]), np.array([1, 10]), 2, - np.full((2, 2, 4), [[0, 0, 0, 0], [3, 3, 3, 2]])), + np.full((2, 2, 4), [[1, 0, 0, 0], [2, 3, 3, 2]])), ], ) def test_multinomial_moment(p, n, size, expected): From f39892c9ab816596db877fda4d998cee4c008185 Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" <41898282+github-actions[bot]@users.noreply.github.com> Date: Fri, 26 Nov 2021 08:05:10 +0000 Subject: [PATCH 5/5] Run pre-commit --- pymc/tests/test_distributions_moments.py | 32 ++++++++++++++---------- 1 file changed, 19 insertions(+), 13 deletions(-) diff --git a/pymc/tests/test_distributions_moments.py b/pymc/tests/test_distributions_moments.py index 7681bfcbc0..ad6e234305 100644 --- a/pymc/tests/test_distributions_moments.py +++ b/pymc/tests/test_distributions_moments.py @@ -1111,24 +1111,30 @@ def test_polyagamma_moment(h, z, size, expected): (np.array([0.25, 0.25, 0.25, 0.25]), 1, None, np.array([1, 0, 0, 0])), (np.array([0.3, 0.6, 0.05, 0.05]), 2, None, np.array([1, 1, 0, 0])), (np.array([0.3, 0.6, 0.05, 0.05]), 10, None, np.array([4, 6, 0, 0])), - (np.array([[0.3, 0.6, 0.05, 0.05], [0.25, 0.25, 0.25, 0.25]]), 10, None, np.array([[4, 6, 0, 0], [4, 2, 2, 2]])), ( - np.array([[0.25, 0.25, 0.25, 0.25], [0.26, 0.26, 0.26, 0.22]]), - np.array([1, 10]), - None, - np.array([[1, 0, 0, 0], [2, 3, 3, 2]]) + np.array([[0.3, 0.6, 0.05, 0.05], [0.25, 0.25, 0.25, 0.25]]), + 10, + None, + np.array([[4, 6, 0, 0], [4, 2, 2, 2]]), ), ( - np.array([0.26, 0.26, 0.26, 0.22]), - np.array([1, 10]), - None, - np.array([[1, 0, 0, 0], [2, 3, 3, 2]]) + np.array([[0.25, 0.25, 0.25, 0.25], [0.26, 0.26, 0.26, 0.22]]), + np.array([1, 10]), + None, + np.array([[1, 0, 0, 0], [2, 3, 3, 2]]), ), ( - np.array([[0.25, 0.25, 0.25, 0.25], [0.26, 0.26, 0.26, 0.22]]), - np.array([1, 10]), - 2, - np.full((2, 2, 4), [[1, 0, 0, 0], [2, 3, 3, 2]])), + np.array([0.26, 0.26, 0.26, 0.22]), + np.array([1, 10]), + None, + np.array([[1, 0, 0, 0], [2, 3, 3, 2]]), + ), + ( + np.array([[0.25, 0.25, 0.25, 0.25], [0.26, 0.26, 0.26, 0.22]]), + np.array([1, 10]), + 2, + np.full((2, 2, 4), [[1, 0, 0, 0], [2, 3, 3, 2]]), + ), ], ) def test_multinomial_moment(p, n, size, expected):