Convolution along one axis only
You can use np.apply_along_axis
to apply np.convolve
along the desired axis. Here is an example of applying a boxcar filter to a 2d array:
import numpy as npa = np.arange(10)a = np.vstack((a,a)).Tfilt = np.ones(3)np.apply_along_axis(lambda m: np.convolve(m, filt, mode='full'), axis=0, arr=a)
This is an easy way to generalize many functions that don't have an axis
argument.
np.apply_along_axis
won't really help you, because you're trying to iterate over two arrays. Effectively, you'd have to use a loop, as described here.
Now, loops are fine if your arrays are small, but if N and P are large, then you probably want to use FFT to convolve instead.
However, you need to appropriately zero pad your arrays first, so that your "full" convolution has the expected shape:
M, N, P = 4, 10, 20A = np.random.randn(M, N)B = np.random.randn(M, P)A_ = np.zeros((M, N+P-1), dtype=A.dtype)A_[:, :N] = AB_ = np.zeros((M, N+P-1), dtype=B.dtype)B_[:, :P] = BA_fft = np.fft.fft(A_, axis=1)B_fft = np.fft.fft(B_, axis=1)C_fft = A_fft * B_fftC = np.real(np.fft.ifft(C_fft))# TestC_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)