Request: Consider using a more compatible polya-gamma random sampler

[zoj613]

Description of your feature request

This is a feature request following a comment by @brandonwillard in pymc-devs/symbolic-pymc#107 .

Looking at the source code, it appears that for random sampling, aesara uses a polya-gamma distribution random sampler whose API isn't quite compatible to scipy's.
https://github.com/pymc-devs/aesara/blob/6c5a8fa044aae27a7ce2b6f90ef49f119bf5bdcd/aesara/tensor/random/basic.py#L387-L421

I would like to propose an alternative that has a scipy compatible API, which I think would ease the maintenance burden. polyagamma (repo: https://github.com/zoj613/polya-gamma) uses the same API as scipy/numpy statistical functions and is flexible and performant.

Here is an example usage:

import numpy as np
from polyagamma import polyagamma

o = polyagamma()

# Get a 5 by 10 array of PG(1, 2) variates.
o = polyagamma(z=2, size=(5, 10))

# Pass sequences as input. Numpy's broadcasting rules apply here.
h = [[1, 2, 3, 4, 5], [9, 8, 7, 6, 5]]
o = polyagamma(h, 1)

# Pass an output array
out = np.empty(5)
polyagamma(out=out)
print(out)

# one can choose a sampling method from {devroye, alternate, gamma, saddle}.
# If not given, the default behaviour is a hybrid sampler that picks a method
# based on the parameter values.
o = polyagamma(method="saddle")

# We can also use an existing instance of `numpy.random.Generator` as a parameter.
# This is useful to reproduce samples generated via a given seed.
rng = np.random.default_rng(12345)
o = polyagamma(random_state=rng)

[brandonwillard]

Yes, using a library with a more NumPy-like interface would be much better; however, the Pólya-gamma distribution itself should actually be removed from this repository.

PolyaGammaRV made it in from symbolic-pymc because some changes to RandomVariable were being made in this repository and I didn't want to leave it behind, so I moved and updated it as well. Regardless, the RandomVariable implementations hosted by this repository are supposed to mirror NumPy's basic numpy.random offerings. I don't think we can reasonably support all types of distributions here.

Once we finalize the changes proposed in #4440, we can move the PolyaGammaRV to PyMC3 and consider using the new package there.

Besides that, one important consideration when switching packages is stability. The polyagamma package doesn't appear to have CI set up, or report the package's testing and coverage status. Those simple metrics would help justify the switch.

[zoj613]

Yes, using a library with a more NumPy-like interface would be much better; however, the Pólya-gamma distribution itself should actually be removed from this repository.

PolyaGammaRV made it in from symbolic-pymc because some changes to RandomVariable were being made in this repository and I didn't want to leave it behind, so I moved and updated it as well. Regardless, the RandomVariable implementations hosted by this repository are supposed to mirror NumPy's basic numpy.random offerings. I don't think we can reasonably support all types of distributions here.

Once we finalize the changes proposed in pymc-devs/pymc3#4440, we can move the PolyaGammaRV to PyMC3 and consider using the new package there.

Besides that, one important consideration when switching packages is stability. The polyagamma package doesn't appear to have CI set up, or report the package's testing and coverage status. Those simple metrics would help justify the switch.

Thanks for the detailed response. I was not aware of these project plans. Nonetheless, I did update the project and added some basic CI and test coverage reporting.

[brandonwillard]

This would need to be done in the v4 branch.

[zoj613]

Are you suggesting i submit a PR to said branch?

[brandonwillard]

Yeah, that would be possible now. The package would likely need to be a conditional dependency (i.e. require a try block that attempts the import), and it would probably belong in pymc3.distributions.continuous.

[zoj613]

@brandonwillard see #4531 for a draft. I am not sure how random generators are passed around by the distributions in the continuous.py module. The distribution class in distribution.py uses a rng variable in the __new__ method that doesnt seem to be directly used by any of the subclasses so I am not sure what is the correct way to make use of it in the PolyaGamma class.

[brandonwillard]

At the moment, it's safe to assume that most design questions like that probably haven't been completely worked out, yet.