Re-implemented zero-inflated Poisson to use standard parameterization (pymc-devs#1300)

* Re-implemented zero-inflated Poisson to use standard parameterization, rather than using indicators

* Fixed typo in latent occupancy

* Fix to ZIP test

Author: Chris Fonnesbeck

pymc3/distributions/discrete.py

@@ -432,20 +432,49 @@ def logp(self, value):
 
 
 class ZeroInflatedPoisson(Discrete):
-    def __init__(self, theta, z, *args, **kwargs):
+    R"""
+    Zero-inflated Poisson log-likelihood.
+
+    Often used to model the number of events occurring in a fixed period
+    of time when the times at which events occur are independent.
+
+    .. math:: 
+    
+        f(x \mid \theta, \psi) = \left\{ \begin{array}{l}
+            (1-\psi) + \psi e^{-\theta}, \text{if } x = 0 \\
+            \psi \frac{e^{-\theta}\theta^x}{x!}, \text{if } x=1,2,3,\ldots
+            \end{array} \right.
+
+    ========  ==========================
+    Support   :math:`x \in \mathbb{N}_0`
+    Mean      :math:`\psi\theta`
+    Variance  :math:`\theta + \frac{1-\psi}{\psi}\theta^2`
+    ========  ==========================
+
+    Parameters
+    ----------
+    theta : float
+        Expected number of occurrences during the given interval
+        (theta >= 0).
+    psi : float
+        Expected proportion of Poisson variates (0 < psi < 1)
+
+    """
+    def __init__(self, theta, psi, *args, **kwargs):
         super(ZeroInflatedPoisson, self).__init__(*args, **kwargs)
         self.theta = theta
-        self.z = z
+        self.psi = psi
         self.pois = Poisson.dist(theta)
-        self.const = ConstantDist.dist(0)
         self.mode = self.pois.mode
 
     def random(self, point=None, size=None, repeat=None):
-        theta = draw_values([self.theta], point=point)
-        # To do: Finish me
-        return None
+            theta, psi = draw_values([self.theta, self.psi], point=point)
+            g = generate_samples(stats.poisson.rvs, theta,
+                                    dist_shape=self.shape,
+                                    size=size)
+            return g * (np.random.random(np.squeeze(g.shape)) < psi)
 
     def logp(self, value):
-        return tt.switch(self.z,
-                        self.pois.logp(value),
-                        self.const.logp(value))
+        return tt.switch(value > 0,
+                        tt.log(self.psi) + self.pois.logp(value),
+                        tt.log((1. - self.psi) + self.psi * tt.exp(-self.theta)))

pymc3/examples/latent_occupancy.py

@@ -47,16 +47,13 @@
 with model:
     # Estimated occupancy
 
-    p = Beta('p', 1, 1)
-
-    # Latent variable for occupancy
-    z = Bernoulli('z', p, y.shape)
+    psi = Beta('psi', 1, 1)
 
     # Estimated mean count
     theta = Uniform('theta', 0, 100)
 
     # Poisson likelihood
-    yd = ZeroInflatedPoisson('y', theta, z, observed=y)
+    yd = ZeroInflatedPoisson('y', theta, psi, observed=y)
 
 
 point = model.test_point
@@ -77,11 +74,9 @@ def run(n=5000):
     if n == "short":
         n = 50
     with model:
-        start = {'p': 0.5, 'z': (y > 0).astype(int), 'theta': 5}
-
-        step1 = Metropolis([theta, p])
+        start = {'psi': 0.5, 'z': (y > 0).astype(int), 'theta': 5}
 
-        step2 = BinaryMetropolis([z])
+        step1 = Metropolis([theta, psi])
 
         trace = sample(n, [step1, step2], start)
 

pymc3/tests/test_distributions.py

@@ -380,7 +380,7 @@ def test_constantdist():
             )
 
 def test_zeroinflatedpoisson():
-    checkd(ZeroInflatedPoisson, I, {'theta': Rplus, 'z': Bool})
+    checkd(ZeroInflatedPoisson, Nat, {'theta': Rplus, 'psi': Unit})
 
 def test_mvnormal():
     for n in [1, 2]:

pymc3/tests/test_distributions_random.py

@@ -201,11 +201,8 @@ def test_constant_dist(self):
         self.check(ConstantDist, c=3)
 
     def test_zero_inflated_poisson(self):
-        # To do: implement ZIP random
-        #self.check(ZeroInflatedPoisson)
-        raise SkipTest(
-            'ZeroInflatedPoisson random sampling not yet implemented.')
-
+        self.check(ZeroInflatedPoisson, theta=1, psi=0.3)
+        
     def test_discrete_uniform(self):
         self.check(DiscreteUniform, lower=0., upper=10)
 
@@ -297,9 +294,7 @@ def test_constant_dist(self):
         self.check(ConstantDist, c=3)
 
     def test_zero_inflated_poisson(self):
-        # To do: implement ZIP random
-        raise SkipTest(
-            'ZeroInflatedPoisson random sampling not yet implemented.')
+       self.check(ZeroInflatedPoisson, theta=1, psi=0.3)
 
     def test_discrete_uniform(self):
         self.check(DiscreteUniform, lower=0., upper=10)
@@ -402,10 +397,7 @@ def test_constantDist(self):
         self.check(ConstantDist, c=(self.ones * 3).astype(int))
 
     def test_zero_inflated_poisson(self):
-        # To do: implement ZIP random
-        raise SkipTest(
-            'ZeroInflatedPoisson random sampling not yet implemented.')
-        self.check(ZeroInflatedPoisson, {}, SkipTest)
+        self.check(ZeroInflatedPoisson, theta=self.ones, psi=self.ones/2)
 
     def test_discrete_uniform(self):
         self.check(DiscreteUniform,
@@ -515,10 +507,7 @@ def test_constantDist(self):
         self.check(ConstantDist, c=(self.ones * 3).astype(int))
 
     def test_zero_inflated_poisson(self):
-        # To do: implement ZIP random
-        raise SkipTest(
-            'ZeroInflatedPoisson random sampling not yet implemented.')
-        self.check(ZeroInflatedPoisson, {})
+        self.check(ZeroInflatedPoisson, theta=self.ones*2, psi=self.ones/3)
 
     def test_discrete_uniform(self):
         self.check(DiscreteUniform, lower=self.zeros.astype(int),