Search code examples
pythonoptimizationjitnumbanumpy-ufunc

Can bessel functions (from scipy.special) be used with Numba

I'm trying to optimize evaluation of an integral (scipy.integrate.quad) over a function containing Bessel functions with Numba.

While Numba seems to work well for "common" numpy functions, it throws an error when I attempt to include the Bessel function:

Untyped global name 'jn': cannot determine Numba type of <class 'numpy.ufunc'>

From a bit of Googling, I have found a Jupyter notebook from the Numba repository that discusses making a j0 function (https://github.com/numba/numba/blob/08d5c889491213288be0d5c7d726c4c34221c35b/examples/notebooks/j0%20in%20Numba.ipynb).

The notebook comments that making the function in numba will be fast, yet the timing results they show at the end indicate ~100x slower performance with numba. Am I missing something obvious here?

And more generally, is it possible to benefit from Numba compiling for scipy Bessel functions?

Solution

  • Normally NumPy and SciPy provide very fast implementations. Numba on the other hand auto-generates code based on a Python function.

    So you can't apply numba on anything except Python functions and if you want nopython mode (you want it if you're interested in speed) not even every Python function. Numba supports only a very limited set of functions and types. And these functions are all re-implemented in Numba, it doesn't use the Python or NumPy functions at all even if it looks like it would!

    So you have auto-generated LLVM code vs. highly optimized custom made C/Fortran code. So you shouldn't expect to gain anything (although numba generally performs great for very small arrays compared to NumPy/SciPy functions - but for even medium sized arrays numba will be slower).

    But in case you want to use some (currently unsupported) functions inside a numba jitted function you have to re-implement it yourself. Except when calling it in a long tight-loop it won't be worth the trouble, just use a normal function. Developer time is often far more important than runtime.

    That doesn't mean numba isn't great. It's great for tasks that require many computations/loops and cannot be realized using existing NumPy or SciPy functions.