仅沿一个轴卷积

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)