numpy - 仅沿一个轴进行卷积

Question

我有两个具有相同第一轴尺寸的二维数组。在 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]

有没有快速的方法？

Answer 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参数的函数的简单方法。

Answer 2

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

Answer 3

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

现在，如果您的数组很小，则循环很好，但如果 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)

Answer 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