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

Question

我有一个二维 numpy 数组 input_array 和两个索引列表（x_coords 和 y_coords）。我想为以 x,y 坐标为中心的每个 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]

虽然这很好用，但我确实觉得它有点麻烦。有没有更直接的方法来做到这一点？另外，将来我想将其扩展到 3D 案例；从 nxmxk 数组中切片 nx3x3 子数组。也可以使用 strides，但到目前为止我还不能让它在 3D 中工作

Answer 1

这是一个使用数组广播的方法：

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()绘制的图表：

Answer 2

一个非常直接的解决方案是列表理解和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 不在矩阵的边界上。