n and p parametrization on Zero Inflated Negative Binomial (#5212)

* add n and p parametrization docstring and code

* add n bound to logcdf

Co-authored-by: Farhan Reynaldo <farhanreynaldo@gmail.com>

Author: farhanreynaldo

pymc/distributions/discrete.py

@@ -1629,6 +1629,15 @@ def ZeroInfNegBinom(a, m, psi, x):
     Var       :math:`\psi\mu +  \left (1 + \frac{\mu}{\alpha} + \frac{1-\psi}{\mu} \right)`
     ========  ==========================
 
+    The zero inflated negative binomial distribution can be parametrized
+    either in terms of mu or p, and either in terms of alpha or n.
+    The link between the parametrizations is given by
+
+    .. math::
+
+        \mu &= \frac{n(1-p)}{p} \\
+        \alpha &= n
+
     Parameters
     ----------
     psi: float
@@ -1637,15 +1646,18 @@ def ZeroInfNegBinom(a, m, psi, x):
         Poission distribution parameter (mu > 0).
     alpha: float
         Gamma distribution parameter (alpha > 0).
-
+    p: float
+        Alternative probability of success in each trial (0 < p < 1).
+    n: float
+        Alternative number of target success trials (n > 0)
     """
 
     rv_op = zero_inflated_neg_binomial
 
     @classmethod
-    def dist(cls, psi, mu, alpha, *args, **kwargs):
+    def dist(cls, psi, mu=None, alpha=None, p=None, n=None, *args, **kwargs):
         psi = at.as_tensor_variable(floatX(psi))
-        n, p = NegativeBinomial.get_n_p(mu=mu, alpha=alpha)
+        n, p = NegativeBinomial.get_n_p(mu=mu, alpha=alpha, p=p, n=n)
         n = at.as_tensor_variable(floatX(n))
         p = at.as_tensor_variable(floatX(p))
         return super().dist([psi, n, p], *args, **kwargs)
@@ -1707,6 +1719,7 @@ def logcdf(value, psi, n, p):
             0 <= value,
             0 <= psi,
             psi <= 1,
+            0 < n,
             0 < p,
             p <= 1,
         )

pymc/tests/test_distributions.py

@@ -1790,13 +1790,29 @@ def logcdf_fn(value, psi, mu, alpha):
             logp_fn,
         )
 
+        self.check_logp(
+            ZeroInflatedNegativeBinomial,
+            Nat,
+            {"psi": Unit, "p": Unit, "n": NatSmall},
+            lambda value, psi, p, n: np.log((1 - psi) * sp.nbinom.pmf(0, n, p))
+            if value == 0
+            else np.log(psi * sp.nbinom.pmf(value, n, p)),
+        )
+
         self.check_logcdf(
             ZeroInflatedNegativeBinomial,
             Nat,
             {"psi": Unit, "mu": Rplusbig, "alpha": Rplusbig},
             logcdf_fn,
         )
 
+        self.check_logcdf(
+            ZeroInflatedNegativeBinomial,
+            Nat,
+            {"psi": Unit, "p": Unit, "n": NatSmall},
+            lambda value, psi, p, n: np.log((1 - psi) + psi * sp.nbinom.cdf(value, n, p)),
+        )
+
         self.check_selfconsistency_discrete_logcdf(
             ZeroInflatedNegativeBinomial,
             Nat,

pymc/tests/test_distributions_random.py

@@ -1467,7 +1467,7 @@ def seeded_zero_inflated_binomial_rng_fn(self):
     ]
 
 
-class TestZeroInflatedNegativeBinomial(BaseTestDistribution):
+class TestZeroInflatedNegativeBinomialMuSigma(BaseTestDistribution):
     def zero_inflated_negbinomial_rng_fn(
         self, size, psi, n, p, negbinomial_rng_fct, random_rng_fct
     ):
@@ -1502,6 +1502,14 @@ def seeded_zero_inflated_negbinomial_rng_fn(self):
     ]
 
 
+class TestZeroInflatedNegativeBinomial(BaseTestDistribution):
+    pymc_dist = pm.ZeroInflatedNegativeBinomial
+    pymc_dist_params = {"psi": 0.9, "n": 12, "p": 0.7}
+    expected_rv_op_params = {"psi": 0.9, "n": 12, "p": 0.7}
+    reference_dist_params = {"psi": 0.9, "n": 12, "p": 0.7}
+    tests_to_run = ["check_pymc_params_match_rv_op"]
+
+
 class TestOrderedLogistic(BaseTestDistribution):
     pymc_dist = _OrderedLogistic
     pymc_dist_params = {"eta": 0, "cutpoints": np.array([-2, 0, 2])}