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

Question

我注意到scipy.specialn阶和参数x的Bessel函数在x中jv(n,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=5Bessel函数作为参数x=10,它返回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 中执行此操作：

\n\n

from scipy import special as sp\nsphj = sp.specfun.sphj\n\nimport ctypes\naddr = ctypes.pythonapi.PyCObject_AsVoidPtr(ctypes.py_object(sphj._cpointer))\n

\n\n

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

\n\n

%%cython -c-Ofast -c-fopenmp --link-args=-fopenmp\nfrom cpython.mem cimport PyMem_Malloc, PyMem_Free\nfrom cython.parallel import prange\nimport numpy as np\nimport cython\nfrom cpython cimport PyCObject_AsVoidPtr\nfrom scipy import special\n\nctypedef void (*sphj_ptr) (const int *n, const double *x, \n                            const int *nm, const double *sj, const double *dj) nogil\n\ncdef sphj_ptr _sphj=<sphj_ptr>PyCObject_AsVoidPtr(special.specfun.sphj._cpointer)\n\n\n@cython.wraparound(False)\n@cython.boundscheck(False)\ndef cython_sphj2(int n, double[::1] x):\n    cdef int count = x.shape[0]\n    cdef double * sj = <double *>PyMem_Malloc(count * sizeof(double) * (n + 1))\n    cdef double * dj = <double *>PyMem_Malloc(count * sizeof(double) * (n + 1))\n    cdef int * mn    = <int *>PyMem_Malloc(count * sizeof(int))\n    cdef double[::1] res = np.empty(count)\n    cdef int i\n    if count < 100:\n        for i in range(x.shape[0]):\n            _sphj(&n, &x[i], mn + i, sj + i*(n+1), dj + i*(n+1))\n            res[i] = sj[i*(n+1) + n]    #choose the element you want here        \n    else:\n        for i in prange(count,  nogil=True):\n            _sphj(&n, &x[i], mn + i, sj + i*(n+1), dj + i*(n+1))\n            res[i] = sj[i*(n+1) + n]    #choose the element you want here\n    PyMem_Free(sj)\n    PyMem_Free(dj)\n    PyMem_Free(mn)\n    return res.base\n

\n\n

sphj()为了进行比较，下面是在 forloop 中调用的 Python 函数：

\n\n

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

\n\n

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

\n\n

x = np.linspace(1, 2, 10)\nr1 = cython_sphj2(4, x)\nr2 = python_sphj(4, x)\nassert np.allclose(r1, r2)\n%timeit cython_sphj2(4, x)\n%timeit python_sphj(4, x)\n

\n\n

结果：

\n\n

10000 loops, best of 3: 21.5 \xc2\xb5s per loop\n10000 loops, best of 3: 28.1 \xc2\xb5s per loop\n

\n\n

以下是 100000 个元素的结果：

\n\n

x = np.linspace(1, 2, 100000)\nr1 = cython_sphj2(4, x)\nr2 = python_sphj(4, x)\nassert np.allclose(r1, r2)\n%timeit cython_sphj2(4, x)\n%timeit python_sphj(4, x)\n

\n\n

结果：

\n\n

10 loops, best of 3: 44.7 ms per loop\n1 loops, best of 3: 231 ms per loop\n

\n