在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”选项，则会得到中心部分，但我想要的是第一部分。当然，我可以从完整的向量中切掉我不需要的部分，但如果可能的话，我想摆脱不必要的计算时间。

有人能建议一个更好的循环矢量化方法吗？

- jacob

---

如果数组很短，那么为什么要费心呢？这是你代码的瓶颈吗？ - Fred Foo

@larsmans 是的，这是瓶颈。我知道这可能看起来不像很多，但从原则上讲，我认为加速可能是2的因素，这将是很棒的。此外，如果有人想要使用更大的数组来完成它，那么这也可能会很有趣。 - jacob

“short” 有多短，而且输入是否静态？ - mtrw

@mtrw 抱歉，我应该在问题中说明 - 在10到50之间。每次调用时，其中一个输入会发生变化，另一个保持不变，所以我想它可以被设置为静态的？ - jacob

你可以考虑使用 scipy 提供的卷积工具：https://docs.scipy.org/doc/scipy/reference/generated/scipy.signal.convolve2d.html - hafiz031

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

根据我的机器性能测试，与np.convolve(a,b)[:len(a)]相比，在n=10..50时表现良好。

此外，这似乎是numba的工作。

- jfs

---

是的，我从未使用过Numba，但它似乎是一个很好的选择。 - reptilicus

不错！它给了我大约2倍的加速比。如果我看一下生成的C代码，for循环中仍然有一些开销（检查是否小于0）。这也可以消除吗？ - jacob

好的，我可以设置wraparound=False，但这并不能提高速度。我想这真的是最好的答案了。谢谢！ - jacob

Answer 2

在numpy中，无法使用矢量化数组操作进行卷积。你最好使用np.convolve(a, b, mode='same')，然后裁剪掉不需要的部分。这可能比你上面纯python双重循环快10倍。如果你真的关心速度，还可以使用Cython自己编写，但它可能不会比np.convolve()更快。

- reptilicus

---

你可能是对的，修剪掉是最好的解决方案。对于我的情况，我会执行np.convolve(a,b)[:len(a)]（这是“full”模式，而不是“same”模式）。 - jacob