在numpy中对二维数组进行矢量化移动窗口操作

Question

我正在对一个固定大小的移动窗口在一个二维数组上进行操作。有没有什么高效的方法可以在Python中实现类似于向量化的操作，而不需要使用循环呢？我现在的结构大概是这样的：

 for i in range(1,xmax-1):
     for j in range(1,ymax-1):
        out[i][j] = f(in[i][j],in[i+1][j],in[i-1][j],in[i][j+1],in[i][j-1],...)

在这个问题中，eat留下的评论提到了可能可以对这个操作进行向量化，但没有进一步的细节。你可以查看这个链接了解更多信息：numpy/scipy中的向量化索引/切片？

Answer 1

你可以使用滚动窗口技术，具体的解释可以参考这里、这里和这里，不过这些主要是针对一维数组的，而我们这里讨论的是二维数组。

下面是使用NumPy实现二维滚动窗口的源代码：

# Rolling window for 2D arrays in NumPy
import numpy as np

def rolling_window(a, shape):  # rolling window for 2D array
    s = (a.shape[0] - shape[0] + 1,) + (a.shape[1] - shape[1] + 1,) + shape
    strides = a.strides + a.strides
    return np.lib.stride_tricks.as_strided(a, shape=s, strides=strides)

a = np.array([[0,  1,  2,  3,  4,  5],
              [6,  7,  8,  9, 10,  11],
              [12, 13, 14, 15, 7,   8],
              [18, 19, 20, 21, 13, 14],
              [24, 25, 26, 27, 19, 20],
              [30, 31, 32, 33, 34, 35]], dtype=np.int)
b = np.arange(36, dtype=np.float).reshape(6,6)
present = np.array([[7,8],[13,14],[19,20]], dtype=np.int)
absent  = np.array([[7,8],[42,14],[19,20]], dtype=np.int)

found = np.all(np.all(rolling_window(a, present.shape) == present, axis=2), axis=2)
print(np.transpose(found.nonzero()))
found = np.all(np.all(rolling_window(b, present.shape) == present, axis=2), axis=2)
print(np.transpose(found.nonzero()))
found = np.all(np.all(rolling_window(a, absent.shape) == absent, axis=2), axis=2)
print(np.transpose(found.nonzero()))

在数组a中，数组present出现了两次，分别在位置[1,1]和[2,4]。

更多的例子可以在我的CoLab笔记本中找到，链接是“在NumPy数组上使用滚动窗口而不使用for循环”。

Answer 2

如果你能把这个函数

f(in[i][j],in[i+1][j],in[i-1][j],in[i][j+1],in[i][j-1],…)

表示成一个线性操作符，那你就可以用 scipy 的 signal.convolve2d 函数来实现这个功能。举个例子，假设你有一个 50x50 的数组 A，你想计算一个新的数组 B，其中每个元素 b[ij] 是数组 A 中四个特定元素的平均值，这四个元素分别是 a[i,j], a[(i-1),j], a[i,(j-1)], a[(i-1),(j-1)]。你可以简单地这样做：

A = # your first array
B = numpy.ones((2,2))/4
C = scipy.signal.convolve2d(A,B, 'valid')

在进行卷积操作时，数组 B 会“滑动”过数组 A，逐个相乘对应的元素，然后把结果加起来。由于边缘效应的影响，你在使用结果数组 C 时需要小心。这里，C 的形状是 49x49，因为在 convolve2d 中使用了 'valid' 参数，这样会丢掉第一行和第一列，因为它们受到边缘效应的影响。如果你想要一个 50x50 的数组，而不丢掉任何数据，你可以把这个参数换成 'same'。

编辑：也许如果你能告诉我更多关于你需要的那个函数的信息，我可以更具体地帮助你把它转化成一个用于进行二维卷积的数组。

希望这对你有帮助！