如何实现NymPy卷积的多维版本?

时间:2018-08-13 11:41:24

标签: python numpy

this post之后,我被问到基于MCVE的另一个问题。我的目标是为任意形状的输入数组实现NumPy's convolve。请考虑我在多元幂级数(多元多项式乘法)的Cauchy product上下文中使用了卷积一词。如我所解释的here,诸如signal.convolvendimage.convolvendimage.filters.convolve之类的SciPy函数对我不起作用。

考虑两个非正方形2D NumPy数组AB

D1=np.array([4,5])
D2=np.array([2,3])
A=np.random.randint(10,size=D1)
B=np.random.randint(10,size=D2)

例如:

[[1 4 4 2 7]
 [6 1 7 5 3]
 [1 4 3 4 8]
 [7 5 8 3 3]]
[[2 2 3]
 [5 2 9]]

现在我可以用C=convolve(A,B)来计算conv(A,B,K)的元素了:

def crop(A,D1,D2):
    return A[tuple(slice(D1[i], D2[i]) for i in range(A.ndim))]

def sumall(A):
    sum1=A
    for k in range(A.ndim):
        sum1 = np.sum(sum1,axis=0)
    return sum1

def flipall(A):
    return A[[slice(None, None, -1)] * A.ndim]

def conv(A,B,K):
    D0=np.zeros(K.shape,dtype=K.dtype)
    return sumall(np.multiply(crop(A,np.maximum(D0,np.minimum(A.shape,K-B.shape)) \
                           ,np.minimum(A.shape,K)), \
                      flipall(crop(B,np.maximum(D0,np.minimum(B.shape,K-A.shape)) \
                           ,np.minimum(B.shape,K)))))

下面的示例K=np.array([0,0])+1conve(A,B,K)结果1*2=2K=np.array([1,0])+1结果5*1+2*6=17K=np.array([0,1])+12*4+1*2=10 K=np.array([1,1])+1给出4*5+6*2+1*1+1*2=36

[[2 10 ...]
 [17 36 ...]
...  ]]

现在,如果我知道AB的尺寸,则可以嵌套一些for循环以填充C,但这对我来说不是这种情况。如何使用conv函数以C的形状填充C.shape=A.shape+B.shape-1 ndarray而不使用for循环?

1 个答案:

答案 0 :(得分:2)

scipy.ndimageastropy中存在许多n维卷积函数。让我们看看是否可以使用它们中的任何一个。

首先,我们需要一些数据进行比较。因此,让我们扩展输入空间:

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

并计算在该空间上的卷积:

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

好吧,现在我们知道期望值,让我们看看是否可以让scipyastropy赋予我们相同的值:

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

如您所见,这只是选择正确的规范化和填充的问题,您可以简单地使用这些库中的任何一个。

我建议使用astropy.convolution.convolve_fft,因为它(基于FFT)可能是最快的。