Bessel 函数(来自 scipy.special)可以与 Numba 一起使用吗

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

我正在尝试优化对包含带有 Numba 的贝塞尔函数的函数的积分 (scipy.integrate.quad) 的计算。

虽然 Numba 似乎适用于 "common" numpy 函数,但当我尝试包含 Bessel 函数时它会抛出错误:

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

通过一些谷歌搜索,我从 Numba 存储库中找到了一个 Jupyter notebook,其中讨论了制作 j0 函数 (https://github.com/numba/numba/blob/08d5c889491213288be0d5c7d726c4c34221c35b/examples/notebooks/j0%20in%20Numba.ipynb)。

笔记本评论说在 numba 中创建函数会很快,但他们最后显示的计时结果表明使用 numba 的性能要慢 100 倍。我在这里遗漏了什么明显的东西吗?

更一般地说,是否可以从针对 scipy 贝塞尔函数的 Numba 编译中获益?

通常 NumPy 和 SciPy 提供非常快 的实现。另一方面,Numba 基于 Python 函数自动生成代码。

因此,除了 Python 函数之外,你不能将 numba 应用于任何东西,如果你想要 nopython 模式(如果你对速度感兴趣,你会想要它)甚至不是每个 Python 函数。 Numba 仅支持一组非常有限的函数和类型。而且这些功能都在 Numba 中重新实现,它根本不使用 Python 或 NumPy 功能,即使它看起来会!

所以你有自动生成的 LLVM 代码与高度优化的定制 C/Fortran 代码。所以你不应该期望获得任何东西(尽管与 NumPy/SciPy 函数相比,numba 对于非常小的数组通常表现很好 - 但即使对于中等大小的数组,numba 也会更慢)。

但如果您想在 numba jitted 函数中使用某些(当前不受支持的)函数,您 必须 自己重新实现它。除非在很长的紧循环中调用它,否则不值得麻烦,只需使用普通函数即可。开发时间通常比运行时间重要得多。

这并不意味着 numba 不好。它非常适合需要许多 computations/loops 而 不能 使用现有 NumPy 或 SciPy 函数实现的任务。

这是您要查找的代码:

# Bessel function of order 1 - note that smaller arguments are added first
@nb.jit(nopython = True, nogil = True, cache = False)
def Bessel1(z):
    if z.real <= 8.:
        t = z / 8.
        fz = z * (-0.2666949632 * t**14 + 1.7629415168000002 * t**12 + -5.6392305344 * t**10 + 11.1861160576 * t**8 + -14.1749644604 * t**6 + 10.6608917307 * t**4 + -3.9997296130249995 * t**2 + 0.49999791505)
    else:
        t = 8. / z
        eta = z - 0.75 * cmath.pi
        fz = cmath.sqrt(2 / (cmath.pi * z)) * ((1.9776e-06 * t**6 + -3.48648e-05 * t**4 + 0.0018309904000000001 * t**2 + 1.00000000195) * cmath.cos(eta) - (-6.688e-07 * t**7 + 8.313600000000001e-06 * t**5 + -0.000200241 * t**3 + 0.04687499895 * t) * cmath.sin(eta))
    return fz

虽然我不记得这个特定实现的精度。请注意,我已经为复数完成了此操作,您可能需要也可能不需要。

大约一年前,我从一本旧的计算教科书中提取了这个。这有点像泰勒级数展开。我忘了具体的方法叫什么了。

注意代码中必须先加上"t"的高次幂。这是因为t总是小于1,如果你把一个小的浮点数加到一个大的浮点数上,舍入误差会大很多。