python中的矢量化球形bessel函數？

Question

我注意到，爲了n和自變量x jv(n,x)的scipy.special貝塞爾函數矢量在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的所有階數。所以如果我想爲參數x=10使用Bessel函數，它會返回n = 1,2,3,4,5。它實際上返回JN並且每通導數：

In [21]: sp.sph_jn(5,10) 
Out[21]: 
(array([-0.05440211, 0.07846694, 0.07794219, -0.03949584, -0.10558929, 
     -0.05553451]), 
array([-0.07846694, -0.0700955 , 0.05508428, 0.09374053, 0.0132988 , 
     -0.07226858])) 

爲什麼在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)) 

然後您可以直接調用Fortran函數在用Cython，注意我用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 

比較，這裏是一個for循環調用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

有結合矢量球貝塞爾函數程序到SciPy的爲scipy.special.spherical_x一個pull request，與x = jn, yn, in, kn。有一點點運氣，他們應該把它變成版本0.18.0。

與np.vectorize（即，for-loop）相比的性能改進取決於函數，但可以是數量級。

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階球貝塞爾函數基本上是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)) 

我得到了以下的定時：

%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