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.

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