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

【问题标题】：Vectorize this convolution type loop more efficiently in numpy在 numpy 中更有效地向量化这个卷积类型循环
【发布时间】：2012-09-30 17:51:42
【问题描述】：

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

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 循环相比，向量太长了）。如果我在卷积中使用“相同”选项，我会得到中心部分，但我想要的是第一部分。当然，我可以从整个向量中去掉我不需要的东西，但如果可能的话，我想去掉不必要的计算时间。有人可以建议对循环进行更好的矢量化吗？

【问题讨论】：

如果数组很短，那何必呢？这是您代码中的瓶颈吗？
@larsmans 是的，这是瓶颈。我知道这可能看起来不多，但原则上我认为加速可能是 2 倍，这会很棒。此外，如果有人想用更大的数组来做这件事可能会很有趣。
“短”有多短，输入是静态的吗？
@mtrw sorry 应该指定这个 - 在 10 到 50 之间（有问题的编辑）。每次调用的一个输入都会改变，另一个保持不变，所以我想它可以设为静态？
你可以考虑使用scipy提供的卷积功能：docs.scipy.org/doc/scipy/reference/generated/…

标签： python performance numpy vectorization

【解决方案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 的工作。

【讨论】：

是的，我从来没有玩过 numba，但它似乎是个不错的选择。
不错！它给了我大约 2 倍的加速因子。如果我查看生成的 C 代码，for 循环中仍有一些开销（检查
好的，我可以设置 wraparound=False，但这并不能提高速度。我想这真的是最好的答案。谢谢！

【解决方案2】：

在 numpy 中无法使用矢量化数组操作进行卷积。你最好的选择是使用 np.convolve(a, b, mode='same') 并修剪掉你不需要的东西。这可能比您上面的纯 python 中的双循环快 10 倍。如果您真的关心速度，您也可以使用 Cython 自己动手，但它可能不会比 np.convolve() 快多少。

【讨论】：

你可能是对的，修剪是最好的解决方案。就我而言，我会做 np.convolve(a,b)[:len(a)] （这是“完整”模式，而不是“相同”模式）。