python中的矢量化球面bessel函数？

Question

我注意到，n阶的^{}贝塞尔函数和自变量x jv(n,x)在x中矢量化：

In [14]: import scipy.special as sp In [16]: sp.jv(1, range(3)) # n=1, [x=0,1,2] Out[16]: array([ 0., 0.44005059, 0.57672481])

但是球面贝塞尔函数没有相应的矢量化形式，sp.sph_jn：

In [19]: sp.sph_jn(1,range(3)) 

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-19-ea59d2f45497> in <module>()
----> 1 sp.sph_jn(1,range(3)) #n=1, 3 value array

/home/glue/anaconda/envs/fibersim/lib/python2.7/site-packages/scipy/special/basic.pyc in sph_jn(n, z)
    262     """
    263     if not (isscalar(n) and isscalar(z)):
--> 264         raise ValueError("arguments must be scalars.")
    265     if (n != floor(n)) or (n < 0):
    266         raise ValueError("n must be a non-negative integer.")

ValueError: arguments must be scalars.

此外，球面贝塞尔函数一次计算N的所有阶数。因此，如果我想要n=5贝塞尔函数作为参数x=10，它将返回n=1,2,3,4,5。它实际上在一个过程中返回jn及其派生：

^{pr2}$

为什么API中存在这种不对称性，有人知道有一个库可以返回矢量化的球形贝塞尔函数，或者至少更快（例如在cython中）？在

Answer 1

您可以编写一个cython函数来加速计算，首先要做的是获得fortran函数SPHJ的地址，下面是如何在Python中执行此操作：

from scipy import special as sp
sphj = sp.specfun.sphj

import ctypes
addr = ctypes.pythonapi.PyCObject_AsVoidPtr(ctypes.py_object(sphj._cpointer))

然后可以在Cython中直接调用fortran函数，注意我使用prange()来使用多核来加速计算：

^{pr2}$

比较一下，下面是在forloop中调用sphj()的Python函数：

import numpy as np
def python_sphj(n, x):
    sphj = special.specfun.sphj
    res = np.array([sphj(n, v)[1][n] for v in x])
    return res

以下是10个元素的%timit结果：

x = np.linspace(1, 2, 10)
r1 = cython_sphj2(4, x)
r2 = python_sphj(4, x)
assert np.allclose(r1, r2)
%timeit cython_sphj2(4, x)
%timeit python_sphj(4, x)

结果是：

10000 loops, best of 3: 21.5 µs per loop
10000 loops, best of 3: 28.1 µs per loop

以下是100000个元素的结果：

x = np.linspace(1, 2, 100000)
r1 = cython_sphj2(4, x)
r2 = python_sphj(4, x)
assert np.allclose(r1, r2)
%timeit cython_sphj2(4, x)
%timeit python_sphj(4, x)

结果是：

10 loops, best of 3: 44.7 ms per loop
1 loops, best of 3: 231 ms per loop

Answer 2

有一个pull request将矢量化球面贝塞尔函数例程合并到SciPy as scipy.special.spherical_x，其中x = jn, yn, in, kn。如果运气好的话，他们应该会把它放到0.18.0版本中。在

相对于np.vectorize（即for循环）的性能改进取决于函数，但可以是数量级的。在

import numpy as np
from scipy import special

@np.vectorize
def sphj_vectorize(n, z):
    return special.sph_jn(n, z)[0][-1]

x = np.linspace(1, 2, 10**5)

%timeit sphj_vectorize(4, x)
1 loops, best of 3: 1.47 s per loop

%timeit special.spherical_jn(4, x)
100 loops, best of 3: 8.07 ms per loop

Answer 3

如果有人还感兴趣的话，我找到了一个比特德·普利克（Ted Pudlik）快17倍的解决方案。我使用的事实是，n阶球面贝塞尔函数本质上是n阶标准贝塞尔函数的1/sqrt（x）乘以已经矢量化的n+1/2阶标准贝塞尔函数：

import numpy as np
from scipy import special

sphj_bessel = lambda n, z: special.jv(n+1/2,z)*np.sqrt(np.pi/2)/(np.sqrt(z))

我得到了以下时间安排：

^{pr2}$