More distributions (#1039)

* Add Weibull and Betaproportion

* Add NegativeBinomials

* Add ZeroInflatedNegativeBinomial

* Fix some unit tests and bugs

* Reviews

* Simplify Weibull

* Fix Weibull constraint

* Improve computation of log_prob for NBLogits

* Add dispatch method for NegativeBinomial; fix unit tests

* Add distributions to doc

Author: dirmeier

docs/source/distributions.rst

@@ -77,6 +77,15 @@ Beta
     :show-inheritance:
     :member-order: bysource
 
+BetaProportion
+--------------
+.. autoclass:: numpyro.distributions.continuous.BetaProportion
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+
 Cauchy
 ------
 .. autoclass:: numpyro.distributions.continuous.Cauchy
@@ -253,6 +262,14 @@ Uniform
     :show-inheritance:
     :member-order: bysource
 
+Weibull
+-------
+.. autoclass:: numpyro.distributions.continuous.Weibull
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
 
 Discrete Distributions
 ======================
@@ -389,6 +406,34 @@ OrderedLogistic
     :show-inheritance:
     :member-order: bysource
 
+NegativeBinomial
+----------------
+.. autofunction:: numpyro.distributions.conjugate.NegativeBinomial
+
+NegativeBinomialLogits
+----------------------
+.. autoclass:: numpyro.distributions.conjugate.NegativeBinomialLogits
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+NegativeBinomialProbs
+---------------------
+.. autoclass:: numpyro.distributions.conjugate.NegativeBinomialProbs
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
+NegativeBinomial2
+-----------------
+.. autoclass:: numpyro.distributions.conjugate.NegativeBinomial2
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
+
 Poisson
 -------
 .. autoclass:: numpyro.distributions.discrete.Poisson
@@ -417,6 +462,10 @@ ZeroInflatedPoisson
     :show-inheritance:
     :member-order: bysource
 
+ZeroInflatedNegativeBinomial2
+-----------------------------
+.. autofunction:: numpyro.distributions.conjugate.ZeroInflatedNegativeBinomial2
+
 
 Directional Distributions
 =========================

numpyro/distributions/__init__.py

@@ -5,10 +5,15 @@
     BetaBinomial,
     DirichletMultinomial,
     GammaPoisson,
+    NegativeBinomial2,
+    NegativeBinomialLogits,
+    NegativeBinomialProbs,
+    ZeroInflatedNegativeBinomial2,
 )
 from numpyro.distributions.continuous import (
     LKJ,
     Beta,
+    BetaProportion,
     Cauchy,
     Chi2,
     Dirichlet,
@@ -30,6 +35,7 @@
     SoftLaplace,
     StudentT,
     Uniform,
+    Weibull,
 )
 from numpyro.distributions.directional import ProjectedNormal, VonMises
 from numpyro.distributions.discrete import (
@@ -88,6 +94,7 @@
     "BernoulliProbs",
     "Beta",
     "BetaBinomial",
+    "BetaProportion",
     "Binomial",
     "BinomialLogits",
     "BinomialProbs",
@@ -127,6 +134,9 @@
     "MultivariateNormal",
     "LowRankMultivariateNormal",
     "Normal",
+    "NegativeBinomialProbs",
+    "NegativeBinomialLogits",
+    "NegativeBinomial2",
     "OrderedLogistic",
     "Pareto",
     "Poisson",
@@ -144,6 +154,8 @@
     "Uniform",
     "Unit",
     "VonMises",
+    "Weibull",
     "ZeroInflatedDistribution",
     "ZeroInflatedPoisson",
+    "ZeroInflatedNegativeBinomial2",
 ]

numpyro/distributions/conjugate.py

@@ -1,13 +1,18 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
 
-from jax import lax, random
+from jax import lax, nn, random
 import jax.numpy as jnp
-from jax.scipy.special import betaln, gammaln
+from jax.scipy.special import betainc, betaln, gammaln
 
 from numpyro.distributions import constraints
 from numpyro.distributions.continuous import Beta, Dirichlet, Gamma
-from numpyro.distributions.discrete import BinomialProbs, MultinomialProbs, Poisson
+from numpyro.distributions.discrete import (
+    BinomialProbs,
+    MultinomialProbs,
+    Poisson,
+    ZeroInflatedDistribution,
+)
 from numpyro.distributions.distribution import Distribution
 from numpyro.distributions.util import is_prng_key, promote_shapes, validate_sample
 
@@ -209,3 +214,79 @@ def mean(self):
     @property
     def variance(self):
         return self.concentration / jnp.square(self.rate) * (1 + self.rate)
+
+    def cdf(self, value):
+        bt = betainc(self.concentration, value + 1.0, self.rate / (self.rate + 1.0))
+        return bt
+
+
+def NegativeBinomial(total_count, probs=None, logits=None, validate_args=None):
+    if probs is not None:
+        return NegativeBinomialProbs(total_count, probs, validate_args=validate_args)
+    elif logits is not None:
+        return NegativeBinomialLogits(total_count, logits, validate_args=validate_args)
+    else:
+        raise ValueError("One of `probs` or `logits` must be specified.")
+
+
+class NegativeBinomialProbs(GammaPoisson):
+    arg_constraints = {
+        "total_count": constraints.positive,
+        "probs": constraints.unit_interval,
+    }
+    support = constraints.nonnegative_integer
+
+    def __init__(self, total_count, probs, validate_args=None):
+        self.total_count, self.probs = promote_shapes(total_count, probs)
+        concentration = total_count
+        rate = 1.0 / probs - 1.0
+        super().__init__(concentration, rate, validate_args=validate_args)
+
+
+class NegativeBinomialLogits(GammaPoisson):
+    arg_constraints = {
+        "total_count": constraints.positive,
+        "logits": constraints.real,
+    }
+    support = constraints.nonnegative_integer
+
+    def __init__(self, total_count, logits, validate_args=None):
+        self.total_count, self.logits = promote_shapes(total_count, logits)
+        concentration = total_count
+        rate = jnp.exp(-logits)
+        super().__init__(concentration, rate, validate_args=validate_args)
+
+    @validate_sample
+    def log_prob(self, value):
+        return -(
+            self.total_count * nn.softplus(self.logits)
+            + value * nn.softplus(-self.logits)
+            + _log_beta_1(self.total_count, value)
+        )
+
+
+class NegativeBinomial2(GammaPoisson):
+    """
+    Another parameterization of GammaPoisson with `rate` is replaced by `mean`.
+    """
+
+    arg_constraints = {
+        "mean": constraints.positive,
+        "concentration": constraints.positive,
+    }
+    support = constraints.nonnegative_integer
+
+    def __init__(self, mean, concentration, validate_args=None):
+        rate = concentration / mean
+        super().__init__(concentration, rate, validate_args=validate_args)
+
+
+def ZeroInflatedNegativeBinomial2(
+    mean, concentration, *, gate=None, gate_logits=None, validate_args=None
+):
+    return ZeroInflatedDistribution(
+        NegativeBinomial2(mean, concentration, validate_args=validate_args),
+        gate=gate,
+        gate_logits=gate_logits,
+        validate_args=validate_args,
+    )

numpyro/distributions/continuous.py

@@ -1426,3 +1426,77 @@ def infer_shapes(low=(), high=()):
         batch_shape = lax.broadcast_shapes(low, high)
         event_shape = ()
         return batch_shape, event_shape
+
+
+class Weibull(Distribution):
+    arg_constraints = {
+        "scale": constraints.positive,
+        "concentration": constraints.positive,
+    }
+    support = constraints.positive
+    reparametrized_params = ["scale", "concentration"]
+
+    def __init__(self, scale, concentration, validate_args=None):
+        self.concentration, self.scale = promote_shapes(concentration, scale)
+        batch_shape = lax.broadcast_shapes(jnp.shape(concentration), jnp.shape(scale))
+        super().__init__(batch_shape=batch_shape, validate_args=validate_args)
+
+    def sample(self, key, sample_shape=()):
+        assert is_prng_key(key)
+        return random.weibull_min(
+            key,
+            scale=self.scale,
+            concentration=self.concentration,
+            shape=sample_shape + self.batch_shape,
+        )
+
+    @validate_sample
+    def log_prob(self, value):
+        ll = -jnp.power(value / self.scale, self.concentration)
+        ll += jnp.log(self.concentration)
+        ll += (self.concentration - 1.0) * jnp.log(value)
+        ll -= self.concentration * jnp.log(self.scale)
+        return ll
+
+    def cdf(self, value):
+        return 1 - jnp.exp(-((value / self.scale) ** self.concentration))
+
+    @property
+    def mean(self):
+        return self.scale * jnp.exp(gammaln(1.0 + 1.0 / self.concentration))
+
+    @property
+    def variance(self):
+        return self.scale ** 2 * (
+            jnp.exp(gammaln(1.0 + 2.0 / self.concentration))
+            - jnp.exp(gammaln(1.0 + 1.0 / self.concentration)) ** 2
+        )
+
+
+class BetaProportion(Beta):
+    """
+    The BetaProportion distribution is a reparameterization of the conventional
+    Beta distribution in terms of a the variate mean and a
+    precision parameter.
+
+    **Reference:**
+     `Beta regression for modelling rates and proportion`, Ferrari Silvia, and
+      Francisco Cribari-Neto. Journal of Applied Statistics  31.7 (2004): 799-815.
+    """
+
+    arg_constraints = {
+        "mean": constraints.unit_interval,
+        "concentration": constraints.positive,
+    }
+    reparametrized_params = ["mean", "concentration"]
+    support = constraints.unit_interval
+
+    def __init__(self, mean, concentration, validate_args=None):
+        self.concentration = jnp.broadcast_to(
+            concentration, lax.broadcast_shapes(jnp.shape(concentration))
+        )
+        super().__init__(
+            mean * concentration,
+            (1.0 - mean) * concentration,
+            validate_args=validate_args,
+        )

test/test_distributions.py

@@ -113,6 +113,7 @@ def __init__(self, rate, *, validate_args=None):
     dist.BernoulliProbs: lambda probs: osp.bernoulli(p=probs),
     dist.BernoulliLogits: lambda logits: osp.bernoulli(p=_to_probs_bernoulli(logits)),
     dist.Beta: lambda con1, con0: osp.beta(con1, con0),
+    dist.BetaProportion: lambda mu, kappa: osp.beta(mu * kappa, (1 - mu) * kappa),
     dist.BinomialProbs: lambda probs, total_count: osp.binom(n=total_count, p=probs),
     dist.BinomialLogits: lambda logits, total_count: osp.binom(
         n=total_count, p=_to_probs_bernoulli(logits)
@@ -149,6 +150,10 @@ def __init__(self, rate, *, validate_args=None):
     dist.VonMises: lambda loc, conc: osp.vonmises(
         loc=np.array(loc, dtype=np.float64), kappa=np.array(conc, dtype=np.float64)
     ),
+    dist.Weibull: lambda scale, conc: osp.weibull_min(
+        c=conc,
+        scale=scale,
+    ),
     _TruncatedNormal: _truncnorm_to_scipy,
 }
 
@@ -164,6 +169,9 @@ def get_sp_dist(jax_dist):
     T(dist.Beta, 0.2, 1.1),
     T(dist.Beta, 1.0, jnp.array([2.0, 2.0])),
     T(dist.Beta, 1.0, jnp.array([[1.0, 1.0], [2.0, 2.0]])),
+    T(dist.BetaProportion, 0.2, 10.0),
+    T(dist.BetaProportion, 0.51, jnp.array([2.0, 1.0])),
+    T(dist.BetaProportion, 0.5, jnp.array([[4.0, 4.0], [2.0, 2.0]])),
     T(dist.Chi2, 2.0),
     T(dist.Chi2, jnp.array([0.3, 1.3])),
     T(dist.Cauchy, 0.0, 1.0),
@@ -301,6 +309,9 @@ def get_sp_dist(jax_dist):
     T(dist.Uniform, 0.0, 2.0),
     T(dist.Uniform, 1.0, jnp.array([2.0, 3.0])),
     T(dist.Uniform, jnp.array([0.0, 0.0]), jnp.array([[2.0], [3.0]])),
+    T(dist.Weibull, 0.2, 1.1),
+    T(dist.Weibull, 2.8, jnp.array([2.0, 2.0])),
+    T(dist.Weibull, 1.8, jnp.array([[1.0, 1.0], [2.0, 2.0]])),
 ]
 
 DIRECTIONAL = [
@@ -346,6 +357,25 @@ def get_sp_dist(jax_dist):
     T(dist.MultinomialProbs, jnp.array([0.2, 0.7, 0.1]), 10),
     T(dist.MultinomialProbs, jnp.array([0.2, 0.7, 0.1]), jnp.array([5, 8])),
     T(dist.MultinomialLogits, jnp.array([-1.0, 3.0]), jnp.array([[5], [8]])),
+    T(dist.NegativeBinomialProbs, 10, 0.2),
+    T(dist.NegativeBinomialProbs, 10, jnp.array([0.2, 0.6])),
+    T(dist.NegativeBinomialProbs, jnp.array([4.2, 10.7, 2.1]), 0.2),
+    T(
+        dist.NegativeBinomialProbs,
+        jnp.array([4.2, 10.7, 2.1]),
+        jnp.array([0.2, 0.6, 0.5]),
+    ),
+    T(dist.NegativeBinomialLogits, 10, -2.1),
+    T(dist.NegativeBinomialLogits, 10, jnp.array([-5.2, 2.1])),
+    T(dist.NegativeBinomialLogits, jnp.array([4.2, 10.7, 2.1]), -5.2),
+    T(
+        dist.NegativeBinomialLogits,
+        jnp.array([4.2, 7.7, 2.1]),
+        jnp.array([4.2, 0.7, 2.1]),
+    ),
+    T(dist.NegativeBinomial2, 0.3, 10),
+    T(dist.NegativeBinomial2, jnp.array([10.2, 7, 31]), 10),
+    T(dist.NegativeBinomial2, jnp.array([10.2, 7, 31]), jnp.array([10.2, 20.7, 2.1])),
     T(dist.OrderedLogistic, -2, jnp.array([-10.0, 4.0, 9.0])),
     T(dist.OrderedLogistic, jnp.array([-4, 3, 4, 5]), jnp.array([-1.5])),
     T(dist.Poisson, 2.0),
@@ -631,7 +661,7 @@ def fn(args):
         # finite diff approximation
         expected_grad = (fn_rhs - fn_lhs) / (2.0 * eps)
         assert jnp.shape(actual_grad[i]) == jnp.shape(repara_params[i])
-        assert_allclose(jnp.sum(actual_grad[i]), expected_grad, rtol=0.02)
+        assert_allclose(jnp.sum(actual_grad[i]), expected_grad, rtol=0.02, atol=0.03)
 
 
 @pytest.mark.parametrize(
@@ -699,7 +729,7 @@ def log_likelihood(*params):
 
     expected = log_likelihood(*params)
     actual = jax.jit(log_likelihood)(*params)
-    assert_allclose(actual, expected, atol=1e-5)
+    assert_allclose(actual, expected, atol=2e-5)
 
 
 @pytest.mark.parametrize(
@@ -823,6 +853,8 @@ def test_gof(jax_dist, sp_dist, params):
         pytest.xfail("incorrect submanifold scaling")
 
     num_samples = 10000
+    if "BetaProportion" in jax_dist.__name__:
+        num_samples = 20000
     rng_key = random.PRNGKey(0)
     d = jax_dist(*params)
     samples = d.sample(key=rng_key, sample_shape=(num_samples,))