Implement Chi distribution helper

Author: tvwenger

docs/api_reference.rst

@@ -28,9 +28,10 @@ Distributions
 .. autosummary::
    :toctree: generated/
 
-   GenExtreme
-   GeneralizedPoisson
+   Chi
    DiscreteMarkovChain
+   GeneralizedPoisson
+   GenExtreme
    R2D2M2CP
    histogram_approximation
 

pymc_experimental/distributions/__init__.py

@@ -17,7 +17,7 @@
 Experimental probability distributions for stochastic nodes in PyMC.
 """
 
-from pymc_experimental.distributions.continuous import GenExtreme
+from pymc_experimental.distributions.continuous import Chi, GenExtreme
 from pymc_experimental.distributions.discrete import GeneralizedPoisson
 from pymc_experimental.distributions.histogram_utils import histogram_approximation
 from pymc_experimental.distributions.multivariate import R2D2M2CP
@@ -29,4 +29,5 @@
     "GenExtreme",
     "R2D2M2CP",
     "histogram_approximation",
+    "Chi",
 ]

pymc_experimental/distributions/continuous.py

@@ -23,6 +23,8 @@
 
 import numpy as np
 import pytensor.tensor as pt
+from pymc import ChiSquared, CustomDist
+from pymc.distributions import transforms
 from pymc.distributions.dist_math import check_parameters
 from pymc.distributions.distribution import Continuous
 from pymc.distributions.shape_utils import rv_size_is_none
@@ -216,3 +218,65 @@ def moment(rv, size, mu, sigma, xi):
         if not rv_size_is_none(size):
             mode = pt.full(size, mode)
         return mode
+
+
+class Chi:
+    r"""
+    :math:`\chi` log-likelihood.
+
+    The pdf of this distribution is
+
+    .. math::
+
+       f(x \mid \nu) = \frac{x^{\nu - 1}e^{-x^2/2}}{2^{\nu/2 - 1}\Gamma(\nu/2)}
+
+    .. 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(0, 10, 200)
+        for df in [1, 2, 3, 6, 9]:
+            pdf = st.chi.pdf(x, df)
+            plt.plot(x, pdf, label=r'$\nu$ = {}'.format(df))
+        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:`\sqrt{2}\frac{\Gamma((\nu + 1)/2)}{\Gamma(\nu/2)}`
+    Variance  :math:`\nu - 2\left(\frac{\Gamma((\nu + 1)/2)}{\Gamma(\nu/2)}\right)^2`
+    ========  =========================================================================
+
+    Parameters
+    ----------
+    nu : tensor_like of float
+        Degrees of freedom (nu > 0).
+
+    Examples
+    --------
+    .. code-block:: python
+        import pymc as pm
+        from pymc_experimental.distributions import Chi
+
+        with pm.Model():
+            x = Chi('x', nu=1)
+    """
+
+    @staticmethod
+    def chi_dist(nu: TensorVariable, size: TensorVariable) -> TensorVariable:
+        return pt.math.sqrt(ChiSquared.dist(nu=nu, size=size))
+
+    def __new__(cls, name, nu, **kwargs):
+        if "observed" not in kwargs:
+            kwargs.setdefault("transform", transforms.log)
+        return CustomDist(name, nu, dist=cls.chi_dist, class_name="Chi", **kwargs)
+
+    @classmethod
+    def dist(cls, nu, **kwargs):
+        return CustomDist.dist(nu, dist=cls.chi_dist, class_name="Chi", **kwargs)

pymc_experimental/tests/distributions/test_continuous.py

@@ -26,6 +26,7 @@
     BaseTestDistributionRandom,
     Domain,
     R,
+    Rplus,
     Rplusbig,
     assert_moment_is_expected,
     check_logcdf,
@@ -35,7 +36,7 @@
 )
 
 # the distributions to be tested
-from pymc_experimental.distributions import GenExtreme
+from pymc_experimental.distributions import Chi, GenExtreme
 
 
 class TestGenExtremeClass:
@@ -149,3 +150,26 @@ class TestGenExtreme(BaseTestDistributionRandom):
         "check_pymc_draws_match_reference",
         "check_rv_size",
     ]
+
+
+class TestChiClass:
+    """
+    Wrapper class so that tests of experimental additions can be dropped into
+    PyMC directly on adoption.
+    """
+
+    def test_logp(self):
+        check_logp(
+            Chi,
+            Rplus,
+            {"nu": Rplus},
+            lambda value, nu: sp.chi.logpdf(value, df=nu),
+        )
+
+    def test_logcdf(self):
+        check_logcdf(
+            Chi,
+            Rplus,
+            {"nu": Rplus},
+            lambda value, nu: sp.chi.logcdf(value, df=nu),
+        )