通过索引列表从numpy数组切片子数组

Question

我有一个二维的 numpy 数组 input_array，还有两个索引列表（x_coords 和 y_coords）。我想从每对 x 和 y 的坐标出发，切出一个 3x3 的子数组。最后的结果会是一个包含多个 3x3 子数组的数组，子数组的数量和我有的坐标对数是一样的。

我希望能避免使用 for 循环。目前我使用的是来自 scipy 食谱的“生命游戏步幅”的一种修改版：http://wiki.scipy.org/Cookbook/GameOfLifeStrides

shape = (input_array.shape[0] - 2, input_array.shape[0] - 2, 3, 3)
strides = input_array.strides + input_array.strides
strided = np.lib.stride_trics.as_strided(input_array, shape=shape, strides=strides).\
    reshape(shape[0]*shape[1], shape[2], shape[3])

这个方法会把原始数组变成一个包含所有可能的 3x3 子数组的（扁平化的）数组。然后我将 x 和 y 的坐标对转换成可以从 strided 中选择我想要的子数组的形式：

coords = x_coords - 1 + (y_coords - 1)*shape[1]
sub_arrays = strided[coords]

虽然这个方法运行得很好，但我觉得有点繁琐。有没有更直接的方法可以做到这一点？另外，将来我还想扩展到三维的情况；从一个 nxmxk 的数组中切出 nx3x3 的子数组。也许使用步幅也可以做到，但到目前为止我还没能在三维中实现这个功能。

Answer 1

一个非常简单的解决办法是使用列表推导和 itertools.product：

import itertools

sub_arrays = [input_array[x-1:x+2, y-1:y+2] 
              for x, y in itertools.product(x_coords, y_coords)]

这个方法会生成所有可能的坐标组合，然后从 input_array 中提取出3x3的数组。不过，这其实有点像一个for循环。而且你需要注意，x_coords和y_coords不能在矩阵的边缘上。

Answer 2

这里有一种使用数组广播的方法：

x = np.random.randint(1, 63, 10)
y = np.random.randint(1, 63, 10)

dy, dx = [grid.astype(int) for grid in np.mgrid[-1:1:3j, -1:1:3j]]
Y = dy[None, :, :] + y[:, None, None]
X = dx[None, :, :] + x[:, None, None]

然后你可以用 a[Y, X] 从 a 中选择特定的块。下面是一个示例代码：

img = np.zeros((64, 64))
img[Y, X] = 1

这里是用 pyplot.imshow() 绘制的图：