Add alternative parameterization to negative binomial distribution #4126 (#4134)

* Add alternative parameters p and n

* Update Release Notes

* Add test

* Add test for invalid initializations

* Refactor tests with pytest.mark.parametrize

* Minor change

Author: ricardoV94

RELEASE-NOTES.md

@@ -15,6 +15,7 @@
 - `sample_posterior_predictive_w` can now feed on `xarray.Dataset` - e.g. from `InferenceData.posterior`. (see [#4042](https://github.com/pymc-devs/pymc3/pull/4042))
 - Add MLDA, a new stepper for multilevel sampling. MLDA can be used when a hierarchy of approximate posteriors of varying accuracy is available, offering improved sampling efficiency especially in high-dimensional problems and/or where gradients are not available (see [#3926](https://github.com/pymc-devs/pymc3/pull/3926))
 - Change SMC metropolis kernel to independent metropolis kernel [#4115](https://github.com/pymc-devs/pymc3/pull/3926))
+- Add alternative parametrization to NegativeBinomial distribution in terms of n and p (see [#4126](https://github.com/pymc-devs/pymc3/issues/4126))
 
 
 ## PyMC3 3.9.3 (11 August 2020)

pymc3/distributions/discrete.py

@@ -612,20 +612,60 @@ def NegBinom(a, m, x):
     Mean      :math:`\mu`
     ========  ==========================
 
+    The 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
     ----------
     mu: float
         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)
     """
 
-    def __init__(self, mu, alpha, *args, **kwargs):
+    def __init__(self, mu=None, alpha=None, p=None, n=None, *args, **kwargs):
         super().__init__(*args, **kwargs)
+        mu, alpha = self.get_mu_alpha(mu, alpha, p, n)
         self.mu = mu = tt.as_tensor_variable(floatX(mu))
         self.alpha = alpha = tt.as_tensor_variable(floatX(alpha))
         self.mode = intX(tt.floor(mu))
 
+    def get_mu_alpha(self, mu=None, alpha=None, p=None, n=None):
+        if alpha is None:
+            if n is not None:
+                alpha = n
+            else:
+                raise ValueError(
+                    "Incompatible parametrization. Must specify either alpha or n."
+                )
+        elif n is not None:
+            raise ValueError(
+                "Incompatible parametrization. Can't specify both alpha and n."
+            )
+
+        if mu is None:
+            if p is not None:
+                mu = alpha * (1 - p) / p
+            else:
+                raise ValueError(
+                    "Incompatible parametrization. Must specify either mu or p."
+                )
+        elif p is not None:
+            raise ValueError(
+                "Incompatible parametrization. Can't specify both mu and p."
+            )
+
+        return mu, alpha
+
     def random(self, point=None, size=None):
         r"""
         Draw random values from NegativeBinomial distribution.

pymc3/tests/test_distributions.py

@@ -795,6 +795,32 @@ def test_fun(value, mu, alpha):
             return sp.nbinom.logpmf(value, alpha, 1 - mu / (mu + alpha))
 
         self.pymc3_matches_scipy(NegativeBinomial, Nat, {"mu": Rplus, "alpha": Rplus}, test_fun)
+        self.pymc3_matches_scipy(
+            NegativeBinomial,
+            Nat,
+            {"p": Unit, "n": Rplus},
+            lambda value, p, n: sp.nbinom.logpmf(value, n, p),
+        )
+
+    @pytest.mark.parametrize(
+        "mu, p, alpha, n, expected",
+        [
+            (5, None, None, None, "Incompatible parametrization. Must specify either alpha or n."),
+            (None, .5, None, None, "Incompatible parametrization. Must specify either alpha or n."),
+            (None, None, None, None, "Incompatible parametrization. Must specify either alpha or n."),
+            (5, None, 2, 2, "Incompatible parametrization. Can't specify both alpha and n."),
+            (None, .5, 2, 2, "Incompatible parametrization. Can't specify both alpha and n."),
+            (None, None, 2, 2, "Incompatible parametrization. Can't specify both alpha and n."),
+            (None, None, 2, None, "Incompatible parametrization. Must specify either mu or p."),
+            (None, None, None, 2, "Incompatible parametrization. Must specify either mu or p."),
+            (5, .5, 2, None, "Incompatible parametrization. Can't specify both mu and p."),
+            (5, .5, None, 2, "Incompatible parametrization. Can't specify both mu and p."),
+        ]
+    )
+    def test_negative_binomial_init_fail(self, mu, p, alpha, n, expected):
+        with Model():
+            with pytest.raises(ValueError, match=expected):
+                NegativeBinomial("x", mu=mu, p=p, alpha=alpha, n=n)
 
     def test_laplace(self):
         self.pymc3_matches_scipy(