Numpy 沿轴卷积 2 个二维数组

Posted 2023-02-14

【中文标题】Numpy 沿轴卷积 2 个二维数组

【英文标题】：Numpy convolving along an axis for 2 2D-arrays

【发布时间】：2021-11-12 02:23:06

【问题描述】：

我有 2 个二维数组。我正在尝试沿轴 1 进行卷积。np.convolve 不提供 axis 参数。答案here 使用np.apply_along_axis 将1 个二维数组与一个一维数组进行卷积。但它不能直接应用于我的用例。问题here 没有答案。

MWE如下。

import numpy as np

a = np.random.randint(0, 5, (2, 5))
"""
a=
array([[4, 2, 0, 4, 3],
       [2, 2, 2, 3, 1]])
"""
b = np.random.randint(0, 5, (2, 2))
"""
b=
array([[4, 3],
       [4, 0]])
"""

# What I want
c = np.convolve(a, b, axis=1)  # axis is not supported as an argument
"""
c=
array([[16, 20,  6, 16, 24,  9],
       [ 8,  8,  8, 12,  4,  0]])
"""

我知道我可以使用np.fft.fft 来做到这一点，但完成一件简单的事情似乎是不必要的步骤。有没有一种简单的方法可以做到这一点？谢谢。

【参考方案1】：

为什么不用zip 做一个列表理解？

>>> np.array([np.convolve(x, y) for x, y in zip(a, b)])
array([[16, 20,  6, 16, 24,  9],
       [ 8,  8,  8, 12,  4,  0]])
>>> 

或者scipy.signal.convolve2d:

>>> from scipy.signal import convolve2d
>>> convolve2d(a, b)[[0, 2]]
array([[16, 20,  6, 16, 24,  9],
       [ 8,  8,  8, 12,  4,  0]])
>>> 

【讨论】：

当有 3 维数组并且需要 1d 卷积时，使用 zip 的列表理解将不起作用。将需要两个循环。 convolve2d 的类似问题。您在 OP 给出的示例中使用了一些技巧，但我认为这是一个有用的问题，通用的答案对社区更有益。

【参考方案2】：

一种可能性是手动去傅里叶谱，然后返回：

n = np.max([a.shape, b.shape]) + 1
np.abs(np.fft.ifft(np.fft.fft(a, n=n) * np.fft.fft(b, n=n))).astype(int)
# array([[16, 20,  6, 16, 24,  9],
#        [ 8,  8,  8, 12,  4,  0]])

【讨论】：

您好我已经提到fft 方法有点迂回。我希望在时域中完成这一切。

请记住，根据信号的长度，走 FFT 路线也可能是最快的。

你能指定速度如何改变信号的长度吗？如果您指出有关此的资源会更好。谢谢！

@NilsWerner 我猜你是在提到当长度是2 的幂时，fft 实现是有效的？但这是一个特例，我想有一个通用的解决方案。 n 的值也应该是 a.shape[1]+b.shape[1]-1。

complexity differences between naive convolution and FFT are well known。而且您始终可以填充信号，使其长度为 2 的幂。

【参考方案3】：

在正交维度上循环是否会被认为太丑陋？除非主要维度非常短，否则这不会增加太多开销。提前创建输出数组可确保无需复制任何内存。

def convolvesecond(a, b):
    N1, L1 = a.shape
    N2, L2 = b.shape
    if N1 != N2:
       raise ValueError("Not compatible")
    c = np.zeros((N1, L1 + L2 - 1), dtype=a.dtype)
    for n in range(N1):
        c[n,:] = np.convolve(a[n,:], b[n,:], 'full')
    return c

对于一般情况（沿一对多维数组的第 k 个轴进行卷积），我会求助于我经常使用的一对辅助函数将多维问题转换为基本的二维情况：

def semiflatten(x, d=0):
    '''SEMIFLATTEN - Permute and reshape an array to convenient matrix form
    y, s = SEMIFLATTEN(x, d) permutes and reshapes the arbitrary array X so 
    that input dimension D (default: 0) becomes the second dimension of the 
    output, and all other dimensions (if any) are combined into the first 
    dimension of the output. The output is always 2-D, even if the input is
    only 1-D.
    If D<0, dimensions are counted from the end.
    Return value S can be used to invert the operation using SEMIUNFLATTEN.
    This is useful to facilitate looping over arrays with unknown shape.'''
    x = np.array(x)
    shp = x.shape
    ndims = x.ndim
    if d<0:
        d = ndims + d
    perm = list(range(ndims))
    perm.pop(d)
    perm.append(d)
    y = np.transpose(x, perm)
    # Y has the original D-th axis last, preceded by the other axes, in order
    rest = np.array(shp, int)[perm[:-1]]
    y = np.reshape(y, [np.prod(rest), y.shape[-1]])
    return y, (d, rest)

def semiunflatten(y, s):
    '''SEMIUNFLATTEN - Reverse the operation of SEMIFLATTEN
    x = SEMIUNFLATTEN(y, s), where Y, S are as returned from SEMIFLATTEN,
    reverses the reshaping and permutation.'''
    d, rest = s
    x = np.reshape(y, np.append(rest, y.shape[-1]))
    perm = list(range(x.ndim))
    perm.pop()
    perm.insert(d, x.ndim-1)
    x = np.transpose(x, perm)
    return x

（注意reshape 和transpose 不创建副本，因此这些函数非常快。）

有了这些，通用形式可以写成：

def convolvealong(a, b, axis=-1):
   a, S1 = semiflatten(a, axis)
   b, S2 = semiflatten(b, axis)
   c = convolvesecond(a, b)
   return semiunflatten(c, S1)