python中的矢量化球形贝塞尔函数？

我注意到n阶贝塞尔函数和参数x

jv（n，x）

在x中是矢量化的：

在[14]中：将scipy.special作为sp导入
在[16]中：sp.jv（1，范围（3））#n=1[x=0,1,2]
Out[16]：数组（[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.

---

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

您可以编写一个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（）

来使用多核加速计算：

%%cython -c-Ofast -c-fopenmp --link-args=-fopenmp
from cpython.mem cimport PyMem_Malloc, PyMem_Free
from cython.parallel import prange
import numpy as np
import cython
from cpython cimport PyCObject_AsVoidPtr
from scipy import special

ctypedef void (*sphj_ptr) (const int *n, const double *x, 
                            const int *nm, const double *sj, const double *dj) nogil

cdef sphj_ptr _sphj=<sphj_ptr>PyCObject_AsVoidPtr(special.specfun.sphj._cpointer)


@cython.wraparound(False)
@cython.boundscheck(False)
def cython_sphj2(int n, double[::1] x):
    cdef int count = x.shape[0]
    cdef double * sj = <double *>PyMem_Malloc(count * sizeof(double) * (n + 1))
    cdef double * dj = <double *>PyMem_Malloc(count * sizeof(double) * (n + 1))
    cdef int * mn    = <int *>PyMem_Malloc(count * sizeof(int))
    cdef double[::1] res = np.empty(count)
    cdef int i
    if count < 100:
        for i in range(x.shape[0]):
            _sphj(&n, &x[i], mn + i, sj + i*(n+1), dj + i*(n+1))
            res[i] = sj[i*(n+1) + n]    #choose the element you want here        
    else:
        for i in prange(count,  nogil=True):
            _sphj(&n, &x[i], mn + i, sj + i*(n+1), dj + i*(n+1))
            res[i] = sj[i*(n+1) + n]    #choose the element you want here
    PyMem_Free(sj)
    PyMem_Free(dj)
    PyMem_Free(mn)
    return res.base

以下是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)

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)

结果是：

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

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

以下是100000个元素的结果：

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)

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)

结果是：

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

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

有一个将向量化的球形贝塞尔函数例程作为

SciPy.special.spherical_x

合并到SciPy中，其中

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

---

如果有人仍然感兴趣，我发现一个解决方案比Ted Pudlik的快17倍。我使用了一个事实，即阶数n的球形贝塞尔函数本质上是阶数n+1/2的标准贝塞尔函数的1/sqrt（x）倍，这已经是矢量化的：

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))

我得到了以下时间安排：

%timeit sphj_vectorize(2, x) # x = np.linspace(1, 2, 10**5)
1 loops, best of 3: 759 ms per loop

%timeit sphj_bessel(2,x) # x = np.linspace(1, 2, 10**5)
10 loops, best of 3: 44.6 ms per loop

---

这是scipy.special中的一个已知功能请求（提供了一些关于如何实现这些功能的矢量化版本的指导）：。如果有人有现成的实现，请发布！谢谢分享