Fix NegativeBinomial parameterization and enable its tests

Author: brandonwillard

pymc3/distributions/discrete.py

@@ -727,35 +727,33 @@ def NegBinom(a, m, x):
 
     @classmethod
     def dist(cls, mu=None, alpha=None, p=None, n=None, *args, **kwargs):
-        mu, alpha = cls.get_mu_alpha(mu, alpha, p, n)
-        mu = aet.as_tensor_variable(floatX(mu))
-        alpha = aet.as_tensor_variable(floatX(alpha))
-        # mode = intX(aet.floor(mu))
-        return super().dist([mu, alpha], *args, **kwargs)
+        n, p = cls.get_mu_alpha(mu, alpha, p, n)
+        n = aet.as_tensor_variable(floatX(n))
+        p = aet.as_tensor_variable(floatX(p))
+        return super().dist([n, p], *args, **kwargs)
 
     @classmethod
     def get_mu_alpha(cls, mu=None, alpha=None, p=None, n=None):
-        if alpha is None:
-            if n is not None:
-                n = aet.as_tensor_variable(intX(n))
-                alpha = n
+        if n is None:
+            if alpha is not None:
+                n = aet.as_tensor_variable(floatX(alpha))
             else:
                 raise ValueError("Incompatible parametrization. Must specify either alpha or n.")
-        elif n is not None:
+        elif alpha is not None:
             raise ValueError("Incompatible parametrization. Can't specify both alpha and n.")
 
-        if mu is None:
-            if p is not None:
-                p = aet.as_tensor_variable(floatX(p))
-                mu = alpha * (1 - p) / p
+        if p is None:
+            if mu is not None:
+                mu = aet.as_tensor_variable(floatX(mu))
+                p = n / (mu + n)
             else:
                 raise ValueError("Incompatible parametrization. Must specify either mu or p.")
-        elif p is not None:
+        elif mu is not None:
             raise ValueError("Incompatible parametrization. Can't specify both mu and p.")
 
-        return mu, alpha
+        return n, p
 
-    def logp(value, mu, alpha):
+    def logp(value, n, p):
         r"""
         Calculate log-probability of NegativeBinomial distribution at specified value.
 
@@ -769,6 +767,8 @@ def logp(value, mu, alpha):
         -------
         TensorVariable
         """
+        alpha = n
+        mu = alpha * (1 - p) / p
         negbinom = bound(
             binomln(value + alpha - 1, value)
             + logpow(mu / (mu + alpha), value)
@@ -779,9 +779,9 @@ def logp(value, mu, alpha):
         )
 
         # Return Poisson when alpha gets very large.
-        return aet.switch(aet.gt(alpha, 1e10), Poisson.dist(mu).logp(value), negbinom)
+        return aet.switch(aet.gt(alpha, 1e10), Poisson.logp(value, mu), negbinom)
 
-    def logcdf(value, mu, alpha):
+    def logcdf(value, n, p):
         """
         Compute the log of the cumulative distribution function for NegativeBinomial distribution
         at the specified value.
@@ -801,20 +801,14 @@ def logcdf(value, mu, alpha):
                 f"NegativeBinomial.logcdf expects a scalar value but received a {np.ndim(value)}-dimensional object."
             )
 
-        # TODO: avoid `p` recomputation if distribution was defined in terms of `p`
-        p = alpha / (mu + alpha)
-
         return bound(
-            aet.log(incomplete_beta(alpha, aet.floor(value) + 1, p)),
+            aet.log(incomplete_beta(n, aet.floor(value) + 1, p)),
             0 <= value,
-            0 < alpha,
+            0 < n,
             0 <= p,
             p <= 1,
         )
 
-    def _distr_parameters_for_repr(self):
-        return self._param_type
-
 
 class Geometric(Discrete):
     R"""

pymc3/tests/test_distributions.py

@@ -1171,7 +1171,6 @@ def modified_scipy_hypergeom_logcdf(value, N, k, n):
             {"N": NatSmall, "k": NatSmall, "n": NatSmall},
         )
 
-    @pytest.mark.xfail(reason="Distribution not refactored yet")
     def test_negative_binomial(self):
         def scipy_mu_alpha_logpmf(value, mu, alpha):
             return sp.nbinom.logpmf(value, alpha, 1 - mu / (mu + alpha))