Numpy convolving along an axis for 2 2D-arrays

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

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

n = np.max([a.shape, b.shape]) + 1np.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]])

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)