Convert zero-inflated distributions to Mixture subclasses

pymc3/distributions/__init__.py

@@ -37,6 +37,7 @@
 from .discrete import Constant
 from .discrete import ZeroInflatedPoisson
 from .discrete import ZeroInflatedNegativeBinomial
+from .discrete import ZeroInflatedBinomial
 from .discrete import DiscreteUniform
 from .discrete import Geometric
 from .discrete import Categorical
@@ -106,6 +107,7 @@
            'Constant',
            'ZeroInflatedPoisson',
            'ZeroInflatedNegativeBinomial',
+           'ZeroInflatedBinomial',
            'DiscreteUniform',
            'Geometric',
            'Categorical',

pymc3/distributions/discrete.py

@@ -11,7 +11,7 @@
 
 __all__ = ['Binomial',  'BetaBinomial',  'Bernoulli',  'DiscreteWeibull',
            'Poisson', 'NegativeBinomial', 'ConstantDist', 'Constant',
-           'ZeroInflatedPoisson', 'ZeroInflatedNegativeBinomial',
+           'ZeroInflatedPoisson', 'ZeroInflatedBinomial', 'ZeroInflatedNegativeBinomial',
            'DiscreteUniform', 'Geometric', 'Categorical']
 
 
@@ -644,6 +644,67 @@ def _repr_latex_(self, name=None, dist=None):
                                                 get_variable_name(psi))
 
 
+class ZeroInflatedBinomial(Discrete):
+    R"""
+    Zero-inflated Binomial log-likelihood.
+
+    .. math::
+
+        f(x \mid n, p, \psi) = \left\{ \begin{array}{l}
+            (1-\psi) + \psi (1-p)^{n}, \text{if } x = 0 \\
+            \psi {n \choose x} p^x (1-p)^{n-x}, \text{if } x=1,2,3,\ldots,n
+            \end{array} \right.
+
+    ========  ==========================
+    Support   :math:`x \in \mathbb{N}_0`
+    Mean      :math:`(1 - \psi) n p`
+    Variance  :math:`(1-\psi) n p [1 - p(1 - \psi n)].`
+    ========  ==========================
+
+    Parameters
+    ----------
+    n : int
+        Number of Bernoulli trials (n >= 0).
+    p : float
+        Probability of success in each trial (0 < p < 1).
+    psi : float
+        Expected proportion of Poisson variates (0 < psi < 1)
+
+    """
+
+    def __init__(self, n, p, psi, *args, **kwargs):
+        super(ZeroInflatedBinomial, self).__init__(*args, **kwargs)
+        self.n = n = tt.as_tensor_variable(n)
+        self.p = p = tt.as_tensor_variable(p)
+        self.psi = psi = tt.as_tensor_variable(psi)
+        self.bin = Binomial.dist(n, p)
+        self.mode = self.bin.mode
+
+    def random(self, point=None, size=None, repeat=None):
+        n, p, psi = draw_values([self.n, self.p, self.psi], point=point)
+        g = generate_samples(stats.binom.rvs, n, p,
+                             dist_shape=self.shape,
+                             size=size)
+        sampled = g * (np.random.random(np.squeeze(g.shape)) < psi)
+        return reshape_sampled(sampled, size, self.shape)
+
+    def logp(self, value):
+        return tt.switch(value > 0,
+                         tt.log(self.psi) + self.bin.logp(value),
+                         tt.log((1. - self.psi) + self.psi * tt.pow(1 - self.p, self.n)))
+
+    def _repr_latex_(self, name=None, dist=None):
+        if dist is None:
+            dist = self
+        n = dist.n
+        p = dist.p
+        psi = dist.psi
+        return r'${} \sim \text{{ZeroInflatedBinomial}}(\mathit{{n}}={}, \mathit{{p}}={}, \mathit{{psi}}={})$'.format(name,
+                                                get_variable_name(n),
+                                                get_variable_name(p),
+                                                get_variable_name(psi))
+
+
 class ZeroInflatedNegativeBinomial(Discrete):
     R"""
     Zero-Inflated Negative binomial log-likelihood.

pymc3/tests/test_distributions.py

@@ -591,6 +591,10 @@ def test_zeroinflatednegativebinomial(self):
         self.checkd(ZeroInflatedNegativeBinomial, Nat,
                     {'mu': Rplusbig, 'alpha': Rplusbig, 'psi': Unit})
 
+    def test_zeroinflatedbinomial(self):
+        self.checkd(ZeroInflatedBinomial, Nat,
+                    {'n': NatSmall, 'p': Unit, 'psi': Unit})
+
     @pytest.mark.parametrize('n', [1, 2, 3])
     def test_mvnormal(self, n):
         self.pymc3_matches_scipy(MvNormal, RealMatrix(5, n),

pymc3/tests/test_distributions_random.py

@@ -338,6 +338,9 @@ class TestZeroInflatedNegativeBinomial(BaseTestCases.BaseTestCase):
     distribution = pm.ZeroInflatedNegativeBinomial
     params = {'mu': 1., 'alpha': 1., 'psi': 0.3}
 
+class TestZeroInflatedBinomial(BaseTestCases.BaseTestCase):
+    distribution = pm.ZeroInflatedBinomial
+    params = {'n': 10, 'p': 0.6, 'psi': 0.3}
 
 class TestDiscreteUniform(BaseTestCases.BaseTestCase):
     distribution = pm.DiscreteUniform