24

I have two 2-D arrays with the same first axis dimensions. In python, I would like to convolve the two matrices along the second axis only. I would like to get C below without computing the convolution along the first axis as well.

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]

Is there a fast way?

Paul
  • 1,670
  • 2
  • 14
  • 20

4 Answers4

21

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 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)

This is an easy way to generalize many functions that don't have an axis argument.

idfah
  • 1,233
  • 9
  • 15
  • 4
    does this have any speedup vs a loop along each row? – endolith Apr 05 '14 at 19:32
  • 1
    @endolith Nope, it doesn't. See [this](http://stackoverflow.com/a/23849233/525169) answer... – Praveen Aug 06 '16 at 00:57
  • `np.apply_along_axis` really can't be used here, since the OP wants it to iterate over two arrays. The way to do that is to simply use a loop, as described [here](http://stackoverflow.com/questions/28898858/python-apply-along-axis-of-multiple-arrays). – Praveen Aug 06 '16 at 01:50
10

With ndimage.convolve1d, you can specify the axis...

Genius
  • 489
  • 1
  • 3
  • 22
Benjamin
  • 11,036
  • 13
  • 67
  • 116
  • 2
    Thanks for pointing that out. The 'weights' argument, however, needs to be 1-D. It is 2-D in my case. – Paul Mar 10 '11 at 19:04
  • @Paul: what is the context? What are the weights, B? – Benjamin Mar 10 '11 at 19:39
  • Each row in A is being filtered by the corresponding row in B. I could implement it like that, just thought there might be a faster way. A is on the order of 10s of gigabytes in size and I use overlap-add. – Paul Mar 11 '11 at 05:50
  • `lfilter` has a similar issue that your can only specify a 1-d filter. I think this is just not a common use case. – Paul Mar 12 '11 at 20:35
  • @Paul - your comment that `Each row in A is filtered by the corresponding row in B` makes me think you want a result of shape `(M,N+P-1)` whose contents are `[[sg.convolve(A[0,:],B[0,:],'full'],[sg.convolve(A[1,:],B[1,:],'full'],...,[sg.convolve(A[M-1,:],B[M-1,:],'full']]`. But your example does something different. Could you explain the desired output again? – mtrw Apr 02 '11 at 20:45
  • This answer still isn't answering the question, as Paul says in his comments. – Praveen Aug 06 '16 at 00:55
5

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, 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)
Community
  • 1
  • 1
Praveen
  • 6,301
  • 3
  • 43
  • 60
3

for 2D arrays, the function scipy.signal.convolve2d is faster and scipy.signal.fftconvolve can be even faster (depending on the dimensions of the arrays):

Here the same code with 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

Answers are all the same with slight differences due to different computation methods used (e.g., fft vs direct multiplication, but i don't know what exaclty convolve2d uses):

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

print(np.max(np.abs(C - C_fft)))
>>>1.84741111298e-13
Chachni
  • 353
  • 3
  • 11