我必须做以下类型的许多循环
for i in range(len(a)):
for j in range(i+1):
c[i] += a[j]*b[i-j]
其中a和b是短阵列(大小相同,大约在10到50之间)。这可以使用卷积有效地完成:
import numpy as np
np.convolve(a, b)
然而,这给了我完整的卷积(即矢量太长,与上面的for循环相比)。如果我在convolve中使用'same'选项,我会得到中心部分,但我想要的是第一部分。当然,我可以从完整的向量中删除我不需要的东西,但是如果可能的话我想摆脱不必要的计算时间。 有人可以提出更好的循环矢量化吗?
答案 0 :(得分:3)
你可以在Cython中编写一个小的C扩展名:
# cython: boundscheck=False
cimport numpy as np
import numpy as np # zeros_like
ctypedef np.float64_t np_t
def convolve_cy_np(np.ndarray[np_t] a not None,
np.ndarray[np_t] b not None,
np.ndarray[np_t] c=None):
if c is None:
c = np.zeros_like(a)
cdef Py_ssize_t i, j, n = c.shape[0]
with nogil:
for i in range(n):
for j in range(i + 1):
c[i] += a[j] * b[i - j]
return c
与我机器上的n=10..50
相比,np.convolve(a,b)[:len(a)]
表现良好。
这似乎也是numba
的工作。
答案 1 :(得分:2)
在numpy中无法使用矢量化数组操作进行卷积。你最好的选择是使用np.convolve(a,b,mode ='same')并修剪掉你不需要的东西。这可能比你上面的纯python中的双循环快10倍。如果你真的关心速度,你也可以使用Cython自己动手,但如果比np.convolve()更快,它可能不会太多。