[英]Vectorized spherical bessel functions in python?
我注意到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
的n=5
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中)?
你可以編寫一個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()
來使用multicore來加速計算:
%%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
比較一下,這是在sphj()
中調用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
如果有人仍然感興趣,我發現一個解決方案比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中作為scipy.special.spherical_x
,其中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
聲明:本站的技術帖子網頁,遵循CC BY-SA 4.0協議,如果您需要轉載,請注明本站網址或者原文地址。任何問題請咨詢:yoyou2525@163.com.