Numpy convolving along an axis for 2 2D-arrays

Question

I have 2 2D-arrays. I am trying to convolve along the axis 1. np.convolve doesn't provide the axis argument. The answer here, convolves 1 2D-array with a 1D array using np.apply_along_axis. But it cannot be directly applied to my use case. The question here doesn't have an answer.

MWE is as follows.

import numpy as np

a = np.random.randint(0, 5, (2, 5))
"""
a=
array([[4, 2, 0, 4, 3],
       [2, 2, 2, 3, 1]])
"""
b = np.random.randint(0, 5, (2, 2))
"""
b=
array([[4, 3],
       [4, 0]])
"""

# What I want
c = np.convolve(a, b, axis=1)  # axis is not supported as an argument
"""
c=
array([[16, 20,  6, 16, 24,  9],
       [ 8,  8,  8, 12,  4,  0]])
"""

I know I can do it using np.fft.fft, but it seems like an unnecessary step to get a simple thing done. Is there a simple way to do this? Thanks.

score 1 · Answer 1 · answered Sep 17 '21 at 08:03

1

Why not just do a list comprehension with zip?

>>> np.array([np.convolve(x, y) for x, y in zip(a, b)])
array([[16, 20,  6, 16, 24,  9],
       [ 8,  8,  8, 12,  4,  0]])
>>>

Or with scipy.signal.convolve2d:

>>> from scipy.signal import convolve2d
>>> convolve2d(a, b)[[0, 2]]
array([[16, 20,  6, 16, 24,  9],
       [ 8,  8,  8, 12,  4,  0]])
>>>

answered Sep 17 '21 at 08:03

U13-Forward

69,221
14
89
114

4

list comprehension with zip won't work when there are 3 dimensional arrays and 1d convolution is needed. Two loops will be needed. Similar problem with `convolve2d`. You're using some hacks for the example the OP has given, but I think this is a useful question and a generic answer would be much more beneficial to the community. – Nagabhushan S N Sep 17 '21 at 08:17

Nils Werner · Answer 2 · 2021-09-17T08:56:46.593

0

One possibility could be to manually go the way to the Fourier spectrum, and back:

n = np.max([a.shape, b.shape]) + 1
np.abs(np.fft.ifft(np.fft.fft(a, n=n) * np.fft.fft(b, n=n))).astype(int)
# array([[16, 20,  6, 16, 24,  9],
#        [ 8,  8,  8, 12,  4,  0]])

edited Sep 17 '21 at 08:56

answered Sep 17 '21 at 08:51

Nils Werner

34,832
7
76
98

2

Hi I have already mentioned `fft` method is a bit roundabout. I was looking to do it all in time-domain. – learner Sep 17 '21 at 10:32
1

Keep in mind that, depending on the length of the signal, going the FFT route may also be the fastest. – Nils Werner Sep 17 '21 at 10:50
Can you please specify how the speed varies length of the signal? Would be better if you point to a resource regarding this. Thanks! – Nagabhushan S N Sep 17 '21 at 11:01
@NilsWerner I guess you are mentioning that `fft` implementation is efficient when the length is a power of `2`? But that is a special case and I would like to have a general solution. Also the value of `n` should be `a.shape[1]+b.shape[1]-1`. – learner Sep 17 '21 at 11:11
1

The [complexity differences between naive convolution and FFT are well known](https://stackoverflow.com/questions/16164749/computational-complexity-of-convolution). And you can always pad your signal so its length is a power of 2. – Nils Werner Sep 17 '21 at 11:22
@learner How is it roundabout? – jtlz2 Dec 20 '21 at 09:28
@jtlz2, what I meant is, converting to another domain, multiplying and then bringing it back seems unnecessary. I wanted to know if everything could be done without domain shift and employing the `axis` argument. – learner Dec 20 '21 at 20:09

score 0 · Answer 3 · answered Oct 28 '21 at 03:06

Would it be considered too ugly to loop over the orthogonal dimension? That would not add much overhead unless the main dimension is very short. Creating the output array ahead of time ensures that no memory needs to be copied about.

def convolvesecond(a, b):
    N1, L1 = a.shape
    N2, L2 = b.shape
    if N1 != N2:
       raise ValueError("Not compatible")
    c = np.zeros((N1, L1 + L2 - 1), dtype=a.dtype)
    for n in range(N1):
        c[n,:] = np.convolve(a[n,:], b[n,:], 'full')
    return c

For the generic case (convolving along the k-th axis of a pair of multidimensional arrays), I would resort to a pair of helper functions I always keep on hand to convert multidimensional problems to the basic 2d case:

def semiflatten(x, d=0):
    '''SEMIFLATTEN - Permute and reshape an array to convenient matrix form
    y, s = SEMIFLATTEN(x, d) permutes and reshapes the arbitrary array X so 
    that input dimension D (default: 0) becomes the second dimension of the 
    output, and all other dimensions (if any) are combined into the first 
    dimension of the output. The output is always 2-D, even if the input is
    only 1-D.
    If D<0, dimensions are counted from the end.
    Return value S can be used to invert the operation using SEMIUNFLATTEN.
    This is useful to facilitate looping over arrays with unknown shape.'''
    x = np.array(x)
    shp = x.shape
    ndims = x.ndim
    if d<0:
        d = ndims + d
    perm = list(range(ndims))
    perm.pop(d)
    perm.append(d)
    y = np.transpose(x, perm)
    # Y has the original D-th axis last, preceded by the other axes, in order
    rest = np.array(shp, int)[perm[:-1]]
    y = np.reshape(y, [np.prod(rest), y.shape[-1]])
    return y, (d, rest)

def semiunflatten(y, s):
    '''SEMIUNFLATTEN - Reverse the operation of SEMIFLATTEN
    x = SEMIUNFLATTEN(y, s), where Y, S are as returned from SEMIFLATTEN,
    reverses the reshaping and permutation.'''
    d, rest = s
    x = np.reshape(y, np.append(rest, y.shape[-1]))
    perm = list(range(x.ndim))
    perm.pop()
    perm.insert(d, x.ndim-1)
    x = np.transpose(x, perm)
    return x

(Note that reshape and transpose do not create copies, so these functions are extremely fast.)

With those, the generic form can be written as:

def convolvealong(a, b, axis=-1):
   a, S1 = semiflatten(a, axis)
   b, S2 = semiflatten(b, axis)
   c = convolvesecond(a, b)
   return semiunflatten(c, S1)

Numpy convolving along an axis for 2 2D-arrays

3 Answers3