added generalized gamma distribution

docs/api_reference.rst

@@ -20,7 +20,7 @@ methods in the current release of PyMC experimental.
 =============================
 
 .. automodule:: pymc_experimental.distributions.histogram_utils
-   :members: histogram_approximation
+   :members: histogram_approximation, GeneralizedGamma
 
 
 :mod:`pymc_experimental.utils`

pymc_experimental/__init__.py

@@ -13,3 +13,6 @@
 
 from pymc_experimental import distributions, gp, utils
 from pymc_experimental.bart import *
+from pymc_experimental.distributions import (
+	GeneralizedGamma
+)

pymc_experimental/distributions/__init__.py

@@ -1,2 +1,5 @@
 from pymc_experimental.distributions import histogram_utils
 from pymc_experimental.distributions.histogram_utils import histogram_approximation
+from pymc_experimental.distributions.continuous import (
+	GeneralizedGamma,
+)

pymc_experimental/distributions/continuous.py

@@ -0,0 +1,175 @@
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import aesara
+import aesara.tensor as at
+from pymc.aesaraf import floatX
+from aesara.tensor.var import TensorConstant, TensorVariable
+from aesara.tensor.random.op import RandomVariable
+
+from aesara.tensor.random.basic import gengamma
+from pymc.distributions.continuous import PositiveContinuous
+from pymc.distributions.shape_utils import rv_size_is_none
+from pymc.distributions.dist_math import check_parameters
+
+
+class GeneralizedGamma(PositiveContinuous):
+    r"""
+    Generalized Gamma log-likelihood.
+
+    The pdf of this distribution is
+
+    .. math::
+
+       f(x \mid \alpha, p, \lambda) =
+        \frac{ p\lambda^{-1} (x/\lambda)^{\alpha - 1} e^{-(x/\lambda)^p}}
+        {\Gamma(\alpha/p)}
+
+    .. plot::
+        :context: close-figs
+
+        import matplotlib.pyplot as plt
+        import numpy as np
+        import scipy.stats as st
+        import arviz as az
+        plt.style.use('arviz-darkgrid')
+        x = np.linspace(1, 50, 1000)
+        alphas = [1,1,2,2]
+        ps = [1, 2, 4, 4]
+        lambds = [10., 10., 10., 20.]
+        for alpha, p, lambd in zip(alphas, ps, lambds):
+            pdf = st.gengamma.pdf(x, alpha/p, p, scale=lambd)
+            plt.plot(x, pdf, label=r'$\alpha$ = {}, $p$ = {}, $\lambda$ = {}'.format(alpha, p, lambd))
+        plt.xlabel('x', fontsize=12)
+        plt.ylabel('f(x)', fontsize=12)
+        plt.legend(loc=1)
+        plt.show()
+
+    ========  ==========================================
+    Support   :math:`x \in [0, \infty)`
+    Mean      :math:`\lambda \frac{\Gamma((\alpha+1)/p)}{\Gamma(\alpha/p)}`
+    Variance  :math:`\lambda^2 \left( \frac{\Gamma((\alpha+2)/p)}{\Gamma(\alpha/p)} - \left(\frac{\Gamma((\alpha+1)/p)}{\Gamma(\alpha/p)}\right)^2 \right)`
+    ========  ==========================================
+
+    Parameters
+    ----------
+    alpha : tensor_like of float, optional
+        Shape parameter :math:`\alpha` (``alpha`` > 0).
+        Defaults to 1.
+    p : tensor_like of float, optional
+        Additional shape parameter `p` (p > 0).
+        Defaults to 1.
+    lambd : tensor_like of float, optional
+        Scale parameter :math:`\lambda` (lambd > 0).
+        Defaults to 1.
+
+    Examples
+    --------
+
+    .. code-block:: python
+        with pm.Model():
+            x = pm.GeneralizedGamma('x', alpha=1, p=2, lambd=5)
+    """
+    rv_op = gengamma
+
+    @classmethod
+    def dist(cls, alpha, p, lambd, **kwargs):
+        alpha = at.as_tensor_variable(floatX(alpha))
+        p = at.as_tensor_variable(floatX(p))
+        lambd = at.as_tensor_variable(floatX(lambd))
+
+        return super().dist([alpha, p, lambd], **kwargs)
+
+    def moment(rv, size, alpha, p, lambd):
+        alpha, p, lambd = at.broadcast_arrays(alpha, p, lambd)
+        moment = lambd * at.gamma((alpha + 1) / p) / at.gamma(alpha / p)
+        if not rv_size_is_none(size):
+            moment = at.full(size, moment)
+        return moment
+
+    def logp(
+        value,
+        alpha: TensorVariable,
+        p: TensorVariable,
+        lambd: TensorVariable,
+    ) -> TensorVariable:
+        """
+        Calculate log-probability of Generalized Gamma distribution at specified value.
+        Parameters
+        ----------
+        value : tensor_like of float
+            Value(s) for which log-probability is calculated. If the log probabilities for multiple
+            values are desired the values must be provided in a numpy array or Aesara tensor.
+        alpha : tensor_like of float
+            Shape parameter (alpha > 0).
+        p : tensor_like of float
+            Shape parameter (p > 0).
+        lambd : tensor_like of float
+            Scale parameter (lambd > 0).
+        Returns
+        -------
+        TensorVariable
+        """
+        logp_expression = (
+            at.log(p)
+            - at.log(lambd)
+            + (alpha - 1) * at.log(value / lambd)
+            - (value / lambd) ** p
+            - at.gammaln(alpha / p)
+        )
+
+        bounded_logp_expression = at.switch(
+            at.gt(value, 0),
+            logp_expression,
+            -np.inf,
+        )
+
+        return check_parameters(
+            bounded_logp_expression,
+            alpha > 0,
+            p > 0,
+            lambd > 0,
+            msg="alpha > 0, p > 0, lambd > 0",
+        )
+
+    def logcdf(
+        value,
+        alpha: TensorVariable,
+        p: TensorVariable,
+        lambd: TensorVariable,
+    ) -> TensorVariable:
+        """
+        Compute the log of the cumulative distribution function for GeneralizedGamma
+        distribution at the specified value.
+        Parameters
+        ----------
+        value : tensor_like of float
+            Value(s) for which log CDF is calculated. If the log CDF for multiple
+            values are desired the values must be provided in a numpy array or Aesara tensor.
+        alpha : tensor_like of float
+            Shape parameter (alpha > 0).
+        p : tensor_like of float
+            Shape parameter (p > 0).
+        lambd : tensor_like of float
+            Scale parameter (lambd > 0).
+        Returns
+        -------
+        TensorVariable
+        """
+        logcdf_expression = at.log(at.gammainc(alpha / p, (value / lambd) ** p))
+
+
+        bounded_logcdf_expression = at.switch(
+            at.gt(value, 0),
+            logcdf_expression,
+            -np.inf,
+        )
+
+        return check_parameters(
+            bounded_logcdf_expression,
+            alpha > 0,
+            p > 0,
+            lambd > 0,
+            msg="alpha > 0, p > 0, lambd > 0",
+        )
+

pymc_experimental/tests/test_distributions.py

@@ -0,0 +1,28 @@
+import pytest
+import aesara
+import scipy.stats
+import scipy.stats.distributions as sp
+from pymc_experimental.distributions import GeneralizedGamma
+from pymc.tests.test_distributions import Rplus, Rplusbig, check_logp, check_logcdf
+
+class TestMatchesScipy:
+
+    def test_generalized_gamma_logp(self):
+        check_logp(
+            GeneralizedGamma,
+            Rplus,
+            {"alpha": Rplusbig, "p": Rplusbig, "lambd": Rplusbig},
+            lambda value, alpha, p, lambd: sp.gengamma.logpdf(value, a=alpha / p, c=p, scale=lambd),
+        )
+
+    @pytest.mark.skipif(
+        condition=(aesara.config.floatX == "float32"),
+        reason="Fails on float32 due to numerical issues",
+    )
+    def test_generalized_gamma_logcdf(self):
+        check_logcdf(
+            GeneralizedGamma,
+            Rplus,
+            {"alpha": Rplusbig, "p": Rplusbig, "lambd": Rplusbig},
+            lambda value, alpha, p, lambd: sp.gengamma.logcdf(value, a=alpha / p, c=p, scale=lambd),
+        )

pymc_experimental/tests/test_distributions_moments.py

@@ -0,0 +1,35 @@
+import pytest
+
+import numpy as np
+import scipy.special as special
+import pymc as pm
+from pymc import Model
+from pymc.tests.test_distributions_moments import assert_moment_is_expected
+from pymc_experimental.distributions import GeneralizedGamma 
+
+
+@pytest.mark.parametrize(
+    "alpha, p, lambd, size, expected",
+    [
+        (1, 1, 2, None, 2),
+        (1, 1, 2, 5, np.full(5, 2)),
+        (1, 1, np.arange(1, 6), None, np.arange(1, 6)),
+        (
+            np.arange(1, 6),
+            2 * np.arange(1, 6),
+            10,
+            (2, 5),
+            np.full(
+                (2, 5),
+                10
+                * special.gamma((np.arange(1, 6) + 1) / (np.arange(1, 6) * 2))
+                / special.gamma(np.arange(1, 6) / (np.arange(1, 6) * 2)),
+            ),
+        ),
+    ],
+)
+def test_generalized_gamma_moment(alpha, p, lambd, size, expected):
+    with Model() as model:
+        GeneralizedGamma("x", alpha=alpha, p=p, lambd=lambd, size=size)
+    assert_moment_is_expected(model, expected)
+

pymc_experimental/tests/test_distributions_random.py

@@ -0,0 +1,22 @@
+import pytest
+
+import numpy as np
+import scipy.special as special
+import pymc_experimental as pmx
+from pymc.tests.test_distributions_random import (
+	BaseTestDistributionRandom, 
+	seeded_scipy_distribution_builder,
+)
+
+
+class TestGeneralizedGamma(BaseTestDistributionRandom):
+    pymc_dist = pmx.GeneralizedGamma
+    pymc_dist_params = {"alpha": 2.0, "p": 3.0, "lambd": 5.0}
+    expected_rv_op_params = {"alpha": 2.0, "p": 3.0, "lambd": 5.0}
+    reference_dist_params = {"a": 2.0 / 3.0, "c": 3.0, "scale": 5.0}
+    reference_dist = seeded_scipy_distribution_builder("gengamma")
+    checks_to_run = [
+        "check_pymc_params_match_rv_op",
+        "check_pymc_draws_match_reference",
+    ]
+