How to implement a multidimensional version of NymPy convolve?

Question

Following this post I was advised to ask a different question based on MCVE. My objective is to implement the NumPy\'s convolve for arbitrary shaped input arrays. Please conside

Answer 1

There exist a multitude of n-dimensional convolution functions in scipy.ndimage and astropy. Let's see if we can use any of them.

First we need some data to compare against. So let's span up the input space:

d0, d1 = np.array(A.shape) + np.array(B.shape) - 1
input_space = np.array(np.meshgrid(np.arange(d0), np.arange(d1))).T.reshape(-1, 2)
# array([[0, 0],
#        [0, 1],
#        [0, 2],
#        [0, 3],
#        [0, 4],
#        [0, 5],
#        [0, 6],
#        [1, 0],
#        [1, 1],
#        ...
#        [4, 5],
#        [4, 6]])

and calculate your convolution over this space:

out = np.zeros((d0, d1))
for K in input_space:
    out[tuple(K)] = conv(A, B, K + 1)

out
# array([[  2.,  10.,  19.,  24.,  30.,  20.,  21.],
#        [ 17.,  36.,  71.,  81., 112.,  53.,  72.],
#        [ 32.,  27., 108.,  74., 121.,  79.,  51.],
#        [ 19.,  46.,  79.,  99., 111.,  67.,  81.],
#        [ 35.,  39., 113.,  76.,  93.,  33.,  27.]])

Okay, now that we know what values to expect, lets see if we can get scipy and astropy to give us the same values:

import scipy.signal
scipy.signal.convolve2d(A, B)  # only 2D!
# array([[  2.,  10.,  19.,  24.,  30.,  20.,  21.],
#        [ 17.,  36.,  71.,  81., 112.,  53.,  72.],
#        [ 32.,  27., 108.,  74., 121.,  79.,  51.],
#        [ 19.,  46.,  79.,  99., 111.,  67.,  81.],
#        [ 35.,  39., 113.,  76.,  93.,  33.,  27.]])

import astropy.convolution
astropy.convolution.convolve_fft(
    np.pad(A, pad_width=((1, 0), (1, 1)), mode='constant'),
    B,
    normalize_kernel=False
)
# array([[  2.,  10.,  19.,  24.,  30.,  20.,  21.],
#        [ 17.,  36.,  71.,  81., 112.,  53.,  72.],
#        [ 32.,  27., 108.,  74., 121.,  79.,  51.],
#        [ 19.,  46.,  79.,  99., 111.,  67.,  81.],
#        [ 35.,  39., 113.,  76.,  93.,  33.,  27.]])
astropy.convolution.convolve(
    np.pad(A, pad_width=((1, 0), (1, 1)), mode='constant'),
    np.pad(B, pad_width=((0, 1), (0, 0)), mode='constant'),
    normalize_kernel=False
)
# array([[  2.,  10.,  19.,  24.,  30.,  20.,  21.],
#        [ 17.,  36.,  71.,  81., 112.,  53.,  72.],
#        [ 32.,  27., 108.,  74., 121.,  79.,  51.],
#        [ 19.,  46.,  79.,  99., 111.,  67.,  81.],
#        [ 35.,  39., 113.,  76.,  93.,  33.,  27.]])

import scipy
scipy.ndimage.filters.convolve(
    np.pad(A, pad_width=((0, 1), (0, 2)), mode='constant'),
    B,
    mode='constant',
    cval=0.0,
    origin=-1
)
# array([[  2.,  10.,  19.,  24.,  30.,  20.,  21.],
#        [ 17.,  36.,  71.,  81., 112.,  53.,  72.],
#        [ 32.,  27., 108.,  74., 121.,  79.,  51.],
#        [ 19.,  46.,  79.,  99., 111.,  67.,  81.],
#        [ 35.,  39., 113.,  76.,  93.,  33.,  27.]])
scipy.ndimage.filters.convolve(
    np.pad(A, pad_width=((1, 0), (1, 1)), mode='constant'),
    B,
    mode='constant',
    cval=0.0
)
# array([[  2.,  10.,  19.,  24.,  30.,  20.,  21.],
#        [ 17.,  36.,  71.,  81., 112.,  53.,  72.],
#        [ 32.,  27., 108.,  74., 121.,  79.,  51.],
#        [ 19.,  46.,  79.,  99., 111.,  67.,  81.],
#        [ 35.,  39., 113.,  76.,  93.,  33.,  27.]])

As you can see, it is just a matter of chosing the right normalization and padding, and you can simply use any of these libraries.

I recommend using astropy.convolution.convolve_fft, as it (being FFT-based) is probably the fastest.