attempt to implement Chi distribution helper

Author: tvwenger

docs/api_reference.rst

@@ -33,6 +33,7 @@ Distributions
    DiscreteMarkovChain
    R2D2M2CP
    histogram_approximation
+   Chi
 
 
 Model Transformations

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,63 @@ def moment(rv, size, mu, sigma, xi):
         if not rv_size_is_none(size):
             mode = pt.full(size, mode)
         return mode
+
+
+def chi_dist(nu: TensorVariable, size: TensorVariable) -> TensorVariable:
+    return pt.math.sqrt(ChiSquared.dist(nu=nu, size=size))
+
+
+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
+
+        with pm.Model():
+            x = pm.Chi('x', nu=1)
+    """
+
+    def __new__(cls, name, nu, **kwargs):
+        if "observed" not in kwargs:
+            kwargs.setdefault("transform", transforms.log)
+        return CustomDist(name, (nu,), dist=chi_dist, class_name="Chi", **kwargs)
+
+    @classmethod
+    def dist(cls, nu, **kwargs):
+        return CustomDist.dist(nu, dist=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,54 @@ 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),
+        )
+
+    """
+    @pytest.mark.parametrize(
+        "nu, size, expected",
+        [
+            (1, None, 1),
+            (1, 5, np.full(5, 1)),
+            (np.arange(1, 6), None, np.arange(1, 6)),
+        ],
+    )
+    def test_chi_moment(self, nu, size, expected):
+        with pm.Model() as model:
+            Chi("x", nu=nu, size=size)
+        assert_moment_is_expected(model, expected)
+    """
+
+
+class TestChi(BaseTestDistributionRandom):
+    pymc_dist = Chi
+    pymc_dist_params = {"nu": 3.0}
+    expected_rv_op_params = {"nu": 3.0}
+    reference_dist_params = {"df": 3.0}
+    reference_dist = seeded_scipy_distribution_builder("chi")
+    tests_to_run = [
+        "check_pymc_params_match_rv_op",
+        "check_pymc_draws_match_reference",
+        "check_rv_size",
+    ]