仅沿一个轴进行卷积答案

【问题标题】：Convolution along one axis only仅沿一个轴进行卷积
【发布时间】：2011-07-10 20:47:50
【问题描述】：

我有两个具有相同第一个轴尺寸的二维数组。在 python 中，我只想沿第二个轴对两个矩阵进行卷积。我想得到C 下面而不计算沿第一个轴的卷积。

import numpy as np
import scipy.signal as sg

M, N, P = 4, 10, 20
A = np.random.randn(M, N)
B = np.random.randn(M, P)

C = sg.convolve(A, B, 'full')[(2*M-1)/2]

有没有快速的方法？

【问题讨论】：

标签： numpy signal-processing scipy linear-algebra convolution

【解决方案1】：

您可以使用np.apply_along_axis 沿所需轴应用np.convolve。这是一个将 boxcar 过滤器应用于 2d 数组的示例：

import numpy as np

a = np.arange(10)
a = np.vstack((a,a)).T

filt = np.ones(3)

np.apply_along_axis(lambda m: np.convolve(m, filt, mode='full'), axis=0, arr=a)

这是概括许多没有axis 参数的函数的简单方法。

【讨论】：

这对每一行的循环是否有任何加速？
@endolith 不，它没有。见this答案...
np.apply_along_axis 真的不能在这里使用，因为 OP 希望它遍历两个数组。这样做的方法是简单地使用循环，如here 所述。

【解决方案2】：

使用ndimage.convolve1d，可以指定轴...

【讨论】：

感谢您指出这一点。然而，“权重”参数必须是一维的。就我而言，它是二维的。
@Paul：上下文是什么？ B 的权重是多少？
A 中的每一行都被 B 中的相应行过滤。我可以这样实现它，只是认为可能有更快的方法。 A 的大小约为 10 GB，我使用重叠相加。
lfilter 有一个类似的问题，您只能指定一维过滤器。我认为这不是一个常见的用例。
@Paul - 你对Each row in A is filtered by the corresponding row in B 的评论让我觉得你想要一个形状为(M,N+P-1) 的结果，其内容是[[sg.convolve(A[0,:],B[0,:],'full'],[sg.convolve(A[1,:],B[1,:],'full'],...,[sg.convolve(A[M-1,:],B[M-1,:],'full']]。但是你的例子做了一些不同的事情。你能再解释一下想要的输出吗？

【解决方案3】：

np.apply_along_axis 不会真正帮助您，因为您正在尝试迭代两个数组。实际上，您必须使用循环，如 here 所述。

现在，如果您的数组很小，则循环很好，但如果 N 和 P 很大，那么您可能希望使用 FFT 来代替。

但是，您需要先对数组进行适当的零填充，以便您的“完整”卷积具有预期的形状：

M, N, P = 4, 10, 20
A = np.random.randn(M, N)
B = np.random.randn(M, P)

A_ = np.zeros((M, N+P-1), dtype=A.dtype)
A_[:, :N] = A
B_ = np.zeros((M, N+P-1), dtype=B.dtype)
B_[:, :P] = B

A_fft = np.fft.fft(A_, axis=1)
B_fft = np.fft.fft(B_, axis=1)
C_fft = A_fft * B_fft

C = np.real(np.fft.ifft(C_fft))

# Test
C_test = np.zeros((M, N+P-1))
for i in range(M):
    C_test[i, :] = np.convolve(A[i, :], B[i, :], 'full')

assert np.allclose(C, C_test)

【讨论】：

【解决方案4】：

对于二维数组，函数 scipy.signal.convolve2d 更快，scipy.signal.fftconvolve 甚至更快（取决于数组的维度）：

这里是相同的代码，N = 100000

import time    
import numpy as np
import scipy.signal as sg

M, N, P = 10, 100000, 20
A = np.random.randn(M, N)
B = np.random.randn(M, P)

T1 = time.time()
C = sg.convolve(A, B, 'full')
print(time.time()-T1)

T1 = time.time()
C_2d = sg.convolve2d(A, B, 'full')
print(time.time()-T1)

T1 = time.time()
C_fft = sg.fftconvolve(A, B, 'full')
print(time.time()-T1)

>>> 12.3
>>> 2.1
>>> 0.6

由于使用的计算方法不同（例如，fft 与直接乘法，但我不知道 exaclty convolve2d 使用什么），答案都是相同的：

print(np.max(np.abs(C - C_2d)))
>>>7.81597009336e-14

print(np.max(np.abs(C - C_fft)))
>>>1.84741111298e-13

【讨论】：