Implement Chi distribution helper

[tvwenger]

Following the discussion here and obsolete PRs here and here, I have attempted to implement a Chi distribution helper class using CustomDist.

The distribution is functionally correct in that it is able to reproduce the scipy chi PDF.

There is an issue with adding this distribution to a pymc model, however. For example:

import pymc as pm
from pymc_experimental.distributions import Chi
with pm.Model() as model:
    pm.Chi("x", df=1)
print(model.initial_point())

results in

TypeError: Cannot convert Type Scalar(float64, shape=()) (of Variable Exp.0) into Type Vector(float64, shape=(1,)). You can try to manually convert Exp.0 into a Vector(float64, shape=(1,)).

I can't quite narrow down this issue, so any help would be appreciated!

[ricardoV94]

Your snippet is usingpm.Chi which shouldn't exist

[tvwenger]

@ricardoV94 Thanks for your help! I've removed those superfluous tests.

Also thanks for catching the issue with my snippet! Here's the correct snippet:

import pymc as pm
from pymc_experimental.distributions import Chi
with pm.Model() as model:
    Chi("x", nu=1)
print(model.initial_point())

results in

Traceback (most recent call last):
  File "/home/twenger/science/python/pymc-experimental/test.py", line 5, in <module>
    print(model.initial_point())
          ^^^^^^^^^^^^^^^^^^^^^
  File "/home/twenger/miniconda3/envs/pymc-experimental-test/lib/python3.11/site-packages/pymc/model/core.py", line 1090, in initial_point
    fn = make_initial_point_fn(model=self, return_transformed=True)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/twenger/miniconda3/envs/pymc-experimental-test/lib/python3.11/site-packages/pymc/initial_point.py", line 140, in make_initial_point_fn
    initial_values = make_initial_point_expression(
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/twenger/miniconda3/envs/pymc-experimental-test/lib/python3.11/site-packages/pymc/initial_point.py", line 284, in make_initial_point_expression
    graph.replace_all(replacements, import_missing=True)
  File "/home/twenger/miniconda3/envs/pymc-experimental-test/lib/python3.11/site-packages/pytensor/graph/fg.py", line 527, in replace_all
    self.replace(var, new_var, **kwargs)
  File "/home/twenger/miniconda3/envs/pymc-experimental-test/lib/python3.11/site-packages/pytensor/graph/fg.py", line 491, in replace
    new_var = var.type.filter_variable(new_var, allow_convert=True)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/twenger/miniconda3/envs/pymc-experimental-test/lib/python3.11/site-packages/pytensor/tensor/type.py", line 272, in filter_variable
    raise TypeError(
TypeError: Cannot convert Type Scalar(float64, shape=()) (of Variable Exp.0) into Type Vector(float64, shape=(1,)). You can try to manually convert Exp.0 into a Vector(float64, shape=(1,)).

[tvwenger]

Alright, I fixed one problem and found another. I was passing the nu parameter to CustomDist as a tuple instead of as parameter. That fixes that error, but introduces a new one for the provided code snippet. The initial logp for the distribution is -inf:

import pymc as pm
from pymc_experimental.distributions import Chi
with pm.Model() as model:
    Chi("x", nu=1)
print(model.initial_point())
# {'x_log__': array(-inf)}

[ricardoV94]

That will be solved by the linked PR which implements initial values for CustomDist.

For now we can provide a maxwell_moment function to the moment kwarg of CustomDist. Probably it's good enough to set the moment to ones. In that case we should reintroduce the moment test that was commented out

[ricardoV94]

Do you also want to add a helper for the Maxwell distribution?

[ricardoV94]

The failing tests should be fixed by #240

[ricardoV94]

Rebased from main, tests should now pass and we can merge