在numpy中更有效地向量化这个卷积类型循环

Question

我必须做以下类型的许多循环

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'选项,我会得到中心部分,但我想要的是第一部分.当然,我可以从完整的向量中删除我不需要的东西,但是如果可能的话我想摆脱不必要的计算时间.有人可以建议更好的循环矢量化吗？

Answer 1

您可以在 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。