jax-ml / jax

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
http://jax.readthedocs.io/
Apache License 2.0
30.51k stars 2.8k forks source link

Updating Scipy Library for Bessel and Spherical Bessel Functions #18119

Open abiros123 opened 1 year ago

abiros123 commented 1 year ago

The scipy subroutines for the spherical Bessel functions and Bessel functions (like scipy.special.spherical_jn) are not included in the JAX library. Although other Bessel functions are included within JAX, the spherical Bessel functions are very important in a variety of applications, such as the Fourier Transform of a function in three-dimensions.

Although the scipy library in JAX does have some modified Bessel function subroutines, expanding it to include the entire Bessel function/Spherical Bessel function subroutines in scipy.special would be extremely helpful.

jakevdp commented 1 year ago

Thanks for the request! We don't have any roadmap for including additional scipy.special funcitonality, but would glady accept contributions.

cgiovanetti commented 1 week ago

We're beginning to pull together some community implementations of various scipy.special functions over in https://github.com/JAXtronomy/spexial as well. We don't have spherical Bessel functions implemented yet, though if you find an implementation we'd also welcome the contribution