Implement Chi distribution helper

.github/workflows/test.yml

@@ -18,15 +18,14 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-latest]
-        floatx: [float64]
         python-version: ["3.9"]
         test-subset:
           - pymc_experimental/tests
       fail-fast: false
     runs-on: ${{ matrix.os }}
     env:
       TEST_SUBSET: ${{ matrix.test-subset }}
-      PYTENSOR_FLAGS: floatX=${{ matrix.floatx }},gcc__cxxflags='-march=native'
+      PYTENSOR_FLAGS: gcc__cxxflags='-march=native'
     defaults:
       run:
         shell: bash -l {0}
@@ -77,21 +76,20 @@ jobs:
         uses: codecov/codecov-action@v2
         with:
           env_vars: TEST_SUBSET
-          name: ${{ matrix.os }} ${{ matrix.floatx }}
+          name: ${{ matrix.os }}
           fail_ci_if_error: false
   windows:
     strategy:
       matrix:
         os: [windows-latest]
-        floatx: [float32]
         python-version: ["3.11"]
         test-subset:
           - pymc_experimental/tests
       fail-fast: false
     runs-on: ${{ matrix.os }}
     env:
       TEST_SUBSET: ${{ matrix.test-subset }}
-      PYTENSOR_FLAGS: floatX=${{ matrix.floatx }},gcc__cxxflags='-march=core2'
+      PYTENSOR_FLAGS: gcc__cxxflags='-march=core2'
     defaults:
       run:
         shell: cmd
@@ -144,5 +142,5 @@ jobs:
         uses: codecov/codecov-action@v2
         with:
           env_vars: TEST_SUBSET
-          name: ${{ matrix.os }} ${{ matrix.floatx }}
+          name: ${{ matrix.os }}
           fail_ci_if_error: false

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

@@ -11,13 +11,10 @@
 #   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 #   See the License for the specific language governing permissions and
 #   limitations under the License.
-import platform
-
 import numpy as np
 import pymc as pm
 
 # general imports
-import pytensor
 import pytest
 import scipy.stats.distributions as sp
 
@@ -26,6 +23,7 @@
     BaseTestDistributionRandom,
     Domain,
     R,
+    Rplus,
     Rplusbig,
     assert_moment_is_expected,
     check_logcdf,
@@ -35,7 +33,7 @@
 )
 
 # the distributions to be tested
-from pymc_experimental.distributions import GenExtreme
+from pymc_experimental.distributions import Chi, GenExtreme
 
 
 class TestGenExtremeClass:
@@ -46,10 +44,6 @@ class TestGenExtremeClass:
     pm.logp(GenExtreme.dist(mu=0.,sigma=1.,xi=0.5),value=-0.01)
     """
 
-    @pytest.mark.xfail(
-        condition=(pytensor.config.floatX == "float32"),
-        reason="PyMC underflows earlier than scipy on float32",
-    )
     def test_logp(self):
         def ref_logp(value, mu, sigma, xi):
             if 1 + xi * (value - mu) / sigma > 0:
@@ -68,13 +62,6 @@ def ref_logp(value, mu, sigma, xi):
             ref_logp,
         )
 
-        if pytensor.config.floatX == "float32":
-            raise Exception("Flaky test: It passed this time, but XPASS is not allowed.")
-
-    @pytest.mark.skipif(
-        (pytensor.config.floatX == "float32" and platform.system() == "Windows"),
-        reason="Scipy gives different results on Windows and does not match with desired accuracy",
-    )
     def test_logcdf(self):
         def ref_logcdf(value, mu, sigma, xi):
             if 1 + xi * (value - mu) / sigma > 0:
@@ -149,3 +136,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),
+        )

pymc_experimental/tests/distributions/test_multivariate.py

@@ -1,6 +1,5 @@
 import numpy as np
 import pymc as pm
-import pytensor
 import pytest
 
 import pymc_experimental as pmx
@@ -96,10 +95,6 @@ def phi_args(self, request, phi_args_base):
             phi_args_base["importance_concentration"] = 10
         return phi_args_base
 
-    @pytest.mark.skipif(
-        pytensor.config.floatX == "float32",
-        reason="pytensor.config.floatX == 'float32', https://github.com/pymc-devs/pymc/issues/6779",
-    )
     def test_init(
         self,
         dims,