bessel_jn for n=0 seems not well implemented

[jecampagne]

Hello,
Although I do not use it correctly, the jax.scipy.special.bessel_jn (jax version in Google collab; 0.4.13) does not match neither the scipy..special.j0 and a JAX version I have cooked from Boost C++ implementation boost/math/special_functions/detail/bessel_j0.hpp

You can have a look a the following notebook on Google Collab.
https://colab.research.google.com/drive/1Cr7oDffnpc6q_hUqNYI423nzjo8MQvn3?usp=sharing

[jakevdp]

Hi - thanks for the report. It looks like bessel_jn uses an iterative approximation that is inaccurate out of some input domain that is not well-documented.

My recommendation would be to remove jax.scipy.special.bessel_jn: it's unclear about what range of values the approximation is suitable for, and it strikes me as somewhat out-of-scope for the JAX package itself. Having an incorrect implementation is arguably worse than having no implementation at all.

What do you think?

[jecampagne]

Hi @jakevdp I agree that "Having an incorrect implementation is arguably worse than having no implementation at all."
Now what is rather strange is that from lax_scipy_test.py

@jtu.sample_product(
    [dict(order=order, z=z, n_iter=n_iter)
     for order, z, n_iter in zip(
         [0, 1, 2, 3, 6], [0.01, 1.1, 11.4, 30.0, 100.6], [5, 20, 50, 80, 200]
     )],
  )
  def testBesselJn(self, order, z, n_iter):
    def lax_fun(z):
      return lsp_special.bessel_jn(z, v=order, n_iter=n_iter)

    def scipy_fun(z):
      vals = [osp_special.jv(v, z) for v in range(order+1)]
      return np.array(vals)

    args_maker = lambda : [z]
    self._CheckAgainstNumpy(scipy_fun, lax_fun, args_maker, rtol=1E-6)
    self._CompileAndCheck(lax_fun, args_maker, rtol=1E-8)

why this test has not trigger an error ???

[jecampagne]

Now you can have a look at the boost C++ implementation here, they motivate 2 kinds of recursive approaches according to the relative strength of the order v and the argument z, while they are specializing the computation for v=0 and v=1. Personally, I have used their j0 rational implementation as shown in the notebook, but I have not an implementation of the computation for any other orders.

Note also that I was looking for Knu and find the TensorFlow Prob implementation, but there is no equivalent for jv-like bessel function.

[jakevdp]

why this test has not trigger an error ???

Because it only tests values for which the implementation converges

[jecampagne]

I have answer on this PR 17038 that the nb I have mentionned in the begenning of this thrad shows how an implementation of j0 and j1 can be donne using Boost algo, and is compliant with the use of jit/vmap/jacrev.