ENH add ZINB distribution and test (#1310)

Author: aloctavodia

pymc3/distributions/__init__.py

@@ -30,6 +30,7 @@
 from .discrete import NegativeBinomial
 from .discrete import ConstantDist
 from .discrete import ZeroInflatedPoisson
+from .discrete import ZeroInflatedNegativeBinomial
 from .discrete import DiscreteUniform
 from .discrete import Geometric
 from .discrete import Categorical
@@ -88,6 +89,7 @@
     'NegativeBinomial',
     'ConstantDist',
     'ZeroInflatedPoisson',
+    'ZeroInflatedNegativeBinomial',
     'DiscreteUniform',
     'Geometric',
     'Categorical',
@@ -109,4 +111,4 @@
     'GARCH11'
 ]
 
-           
+           

pymc3/distributions/discrete.py

@@ -9,8 +9,9 @@
 from .distribution import Discrete, draw_values, generate_samples
 
 __all__ = ['Binomial',  'BetaBinomial',  'Bernoulli',  'Poisson',
-           'NegativeBinomial', 'ConstantDist', 'ZeroInflatedPoisson',
-           'DiscreteUniform', 'Geometric', 'Categorical']
+           'NegativeBinomial', 'ConstantDist', 'ZeroInflatedPoisson', 
+           'ZeroInflatedNegativeBinomial', 'DiscreteUniform', 'Geometric',
+           'Categorical']
 
 
 class Binomial(Discrete):
@@ -478,3 +479,55 @@ def logp(self, 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)))
+
+
+class ZeroInflatedNegativeBinomial(Discrete):
+    R"""
+    Zero-Inflated Negative binomial log-likelihood.
+
+    The Zero-inflated version of the Negative Binomial (NB).
+    The NB distribution describes a Poisson random variable
+    whose rate parameter is gamma distributed. 
+
+    .. math::
+
+       f(x \mid \mu, \alpha, \psi) = \left\{ \begin{array}{l}
+            (1-\psi) + \psi \left (\frac{\alpha}{\alpha+\mu} \right) ^\alpha, \text{if } x = 0 \\
+            \psi \frac{\Gamma(x+\alpha)}{x! \Gamma(\alpha)} \left (\frac{\alpha}{\mu+\alpha} \right)^\alpha \left( \frac{\mu}{\mu+\alpha} \right)^x, \text{if } x=1,2,3,\ldots
+            \end{array} \right.
+
+    ========  ==========================
+    Support   :math:`x \in \mathbb{N}_0`
+    Mean      :math:`\psi\mu`
+    Var       :math:`\psi\mu +  \left (1 + \frac{\mu}{\alpha} + \frac{1-\psi}{\mu} \right)`
+    ========  ==========================
+
+    Parameters
+    ----------
+    mu : float
+        Poission distribution parameter (mu > 0).
+    alpha : float
+        Gamma distribution parameter (alpha > 0).
+    psi : float
+        Expected proportion of NegativeBinomial variates (0 < psi < 1)
+    """
+    def __init__(self, mu, alpha, psi, *args, **kwargs):
+        super(ZeroInflatedNegativeBinomial, self).__init__(*args, **kwargs)
+        self.mu = mu
+        self.alpha = alpha
+        self.psi = psi
+        self.nb = NegativeBinomial.dist(mu, alpha)
+        self.mode = self.nb.mode
+        
+    def random(self, point=None, size=None, repeat=None):
+        mu, alpha, psi = draw_values([self.mu, self.alpha, self.psi], point=point)
+        g = generate_samples(stats.gamma.rvs, alpha, scale=mu/alpha,
+                             dist_shape=self.shape,
+                             size=size)
+        g[g == 0] = np.finfo(float).eps  # Just in case 
+        return stats.poisson.rvs(g) * (np.random.random(np.squeeze(g.shape)) < psi)
+
+    def logp(self, value):
+        return tt.switch(value > 0,
+                        tt.log(self.psi) + self.nb.logp(value),
+                        tt.log((1. - self.psi) + self.psi * (self.alpha/(self.alpha+self.mu))**self.alpha))

pymc3/tests/test_distributions.py

@@ -6,8 +6,8 @@
 from ..model import Model, Point, Potential
 from ..blocking import DictToVarBijection, DictToArrayBijection, ArrayOrdering
 from ..distributions import (DensityDist, Categorical, Multinomial, VonMises, Dirichlet,
-                            MvStudentT, MvNormal, ZeroInflatedPoisson, ConstantDist,
-                            Poisson, Bernoulli, Beta, BetaBinomial, StudentTpos,
+                            MvStudentT, MvNormal, ZeroInflatedPoisson, ZeroInflatedNegativeBinomial,
+                            ConstantDist, Poisson, Bernoulli, Beta, BetaBinomial, StudentTpos,
                             StudentT, Weibull, Pareto, InverseGamma, Gamma, Cauchy,
                             HalfCauchy, Lognormal, Laplace, NegativeBinomial, Geometric,
                             Exponential, ExGaussian, Normal, Flat, LKJCorr, Wald,
@@ -382,6 +382,9 @@ def test_constantdist():
 def test_zeroinflatedpoisson():
     checkd(ZeroInflatedPoisson, Nat, {'theta': Rplus, 'psi': Unit})
 
+def test_zeroinflatednegativebinomial():
+    checkd(ZeroInflatedNegativeBinomial, Nat, {'mu': Rplusbig, 'alpha': Rplusbig, 'psi': Unit})
+
 def test_mvnormal():
     for n in [1, 2]:
         yield check_mvnormal, n

pymc3/tests/test_distributions_random.py

@@ -11,8 +11,8 @@
     I, Simplex, Vector, PdMatrix)
 
 from ..distributions import (DensityDist, Categorical, Multinomial, VonMises, Dirichlet,
-                            MvStudentT, MvNormal, ZeroInflatedPoisson, ConstantDist,
-                            Poisson, Bernoulli, Beta, BetaBinomial, StudentTpos, 
+                            MvStudentT, MvNormal, ZeroInflatedPoisson, ZeroInflatedNegativeBinomial,
+                            ConstantDist, Poisson, Bernoulli, Beta, BetaBinomial, StudentTpos, 
                             StudentT, Weibull, Pareto, InverseGamma, Gamma, Cauchy,
                             HalfCauchy, Lognormal, Laplace, NegativeBinomial, Geometric,
                             Exponential, ExGaussian, Normal, Flat, LKJCorr, Wald, 
@@ -202,7 +202,10 @@ def test_constant_dist(self):
 
     def test_zero_inflated_poisson(self):
         self.check(ZeroInflatedPoisson, theta=1, psi=0.3)
-        
+
+    def test_zero_inflated_negative_binomial(self):
+        self.check(ZeroInflatedNegativeBinomial, mu=1., alpha=1., psi=0.3)
+
     def test_discrete_uniform(self):
         self.check(DiscreteUniform, lower=0., upper=10)
 
@@ -295,6 +298,9 @@ def test_constant_dist(self):
 
     def test_zero_inflated_poisson(self):
        self.check(ZeroInflatedPoisson, theta=1, psi=0.3)
+       
+    def test_zero_inflated_negative_binomial(self):
+        self.check(ZeroInflatedNegativeBinomial, mu=1., alpha=1., psi=0.3)
 
     def test_discrete_uniform(self):
         self.check(DiscreteUniform, lower=0., upper=10)
@@ -398,6 +404,9 @@ def test_constantDist(self):
 
     def test_zero_inflated_poisson(self):
         self.check(ZeroInflatedPoisson, theta=self.ones, psi=self.ones/2)
+        
+    def test_zero_inflated_negative_binomial(self):
+        self.check(ZeroInflatedNegativeBinomial, mu=self.ones, alpha=self.ones, psi=self.ones/2)
 
     def test_discrete_uniform(self):
         self.check(DiscreteUniform,
@@ -508,6 +517,9 @@ def test_constantDist(self):
 
     def test_zero_inflated_poisson(self):
         self.check(ZeroInflatedPoisson, theta=self.ones*2, psi=self.ones/3)
+        
+    def test_zero_inflated_negative_binomial(self):
+        self.check(ZeroInflatedNegativeBinomial, mu=self.ones*2, alpha=self.ones*2, psi=self.ones/3)
 
     def test_discrete_uniform(self):
         self.check(DiscreteUniform, lower=self.zeros.astype(int),