diff --git a/docs/api_reference.rst b/docs/api_reference.rst index 06375fca..1612ac72 100644 --- a/docs/api_reference.rst +++ b/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` diff --git a/pymc_experimental/__init__.py b/pymc_experimental/__init__.py index 8946a3a1..97073314 100644 --- a/pymc_experimental/__init__.py +++ b/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 +) diff --git a/pymc_experimental/distributions/__init__.py b/pymc_experimental/distributions/__init__.py index 41ff5483..a368a847 100644 --- a/pymc_experimental/distributions/__init__.py +++ b/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, +) \ No newline at end of file diff --git a/pymc_experimental/distributions/continuous.py b/pymc_experimental/distributions/continuous.py new file mode 100644 index 00000000..ba36eecb --- /dev/null +++ b/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", + ) + diff --git a/pymc_experimental/tests/test_distributions.py b/pymc_experimental/tests/test_distributions.py new file mode 100644 index 00000000..bcbd3b41 --- /dev/null +++ b/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), + ) \ No newline at end of file diff --git a/pymc_experimental/tests/test_distributions_moments.py b/pymc_experimental/tests/test_distributions_moments.py new file mode 100644 index 00000000..3152aeea --- /dev/null +++ b/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) + diff --git a/pymc_experimental/tests/test_distributions_random.py b/pymc_experimental/tests/test_distributions_random.py new file mode 100644 index 00000000..d1e60a60 --- /dev/null +++ b/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", + ] +