Skip to content

Commit 917e95a

Browse files
sagartomarricardoV94
authored andcommitted
Corrected the mean for ZeroInflatedBinomial
1 parent fe2d101 commit 917e95a

File tree

2 files changed

+6
-6
lines changed

2 files changed

+6
-6
lines changed

pymc/distributions/discrete.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1462,7 +1462,7 @@ class ZeroInflatedBinomial(Discrete):
14621462
14631463
======== ==========================
14641464
Support :math:`x \in \mathbb{N}_0`
1465-
Mean :math:`(1 - \psi) n p`
1465+
Mean :math:`\psi n p`
14661466
Variance :math:`(1-\psi) n p [1 - p(1 - \psi n)].`
14671467
======== ==========================
14681468
@@ -1487,7 +1487,7 @@ def dist(cls, psi, n, p, *args, **kwargs):
14871487
return super().dist([psi, n, p], *args, **kwargs)
14881488

14891489
def get_moment(rv, size, psi, n, p):
1490-
mean = at.round((1 - psi) * n * p)
1490+
mean = at.round(psi * n * p)
14911491
if not rv_size_is_none(size):
14921492
mean = at.full(size, mean)
14931493
return mean

pymc/tests/test_distributions_moments.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -553,11 +553,11 @@ def test_zero_inflated_poisson_moment(psi, theta, size, expected):
553553
@pytest.mark.parametrize(
554554
"psi, n, p, size, expected",
555555
[
556-
(0.2, 7, 0.7, None, 4),
557-
(0.2, 7, 0.3, 5, np.full(5, 2)),
558-
(0.6, 25, np.arange(1, 6) / 10, None, np.arange(1, 6)),
556+
(0.8, 7, 0.7, None, 4),
557+
(0.8, 7, 0.3, 5, np.full(5, 2)),
558+
(0.4, 25, np.arange(1, 6) / 10, None, np.arange(1, 6)),
559559
(
560-
0.6,
560+
0.4,
561561
25,
562562
np.arange(1, 6) / 10,
563563
(2, 5),

0 commit comments

Comments
 (0)