如何从大型2D NumPy数组中提取多组2D切片？

问题 #

现有大型2D NumPy数组，以及对应提取目标的x、y索引数组，提取单个索引对的元素很简单：

import numpy as np

h = np.arange(49).reshape(7,7)
# h = [[0, 1, 2, 3, 4, 5, 6],
#      [7, 8, 9, 10, 11, 12, 13],
#      [14, 15, 16, 17, 18, 19, 20],
#      [21, 22, 23, 24, 25, 26, 27],
#      [28, 29, 30, 31, 32, 33, 34],
#      [35, 36, 37, 38, 39, 40, 41],
#      [42, 43, 44, 45, 46, 47, 48]]

x_indices = np.array([1,3,4])
y_indices = np.array([2,3,5])

reduced_h = h[x_indices, y_indices]
# reduced_h = [ 9, 24, 33]

但需要针对每个x、y索引对，提取以该坐标为中心、向各方向延伸a个索引的正方形切片，返回由这些小2D数组组成的数组。例如当a=1时，返回形状为(len(x_indices),2a+1,2a+1)的数组：

# 预期结果
reduced_h = [[[1,2,3],[8,9,10],[15,16,17]], 
             [[16,17,18],[23,24,25],[30,31,32]], 
             [[25,26,27],[32,33,34],[39,40,41]]]

尝试直接切片h[x_indices-a : x_indices+a, y_indices-a : y_indices+a]失败，用linspace生成索引范围后，h[xrange,yrange]会按元素索引而非切片提取，请问如何实现无循环的广播式切片提取？

解决方案 #

可以通过广播生成所有需要的索引坐标，再利用NumPy的高级索引实现批量提取，无需循环：

步骤1：生成偏移量范围 #

先创建对应正方形切片的行列偏移量，比如a=1时，偏移量是[-1,0,1]：

a = 1
offsets = np.arange(-a, a+1)  # [-1, 0, 1]

步骤2：广播生成所有目标索引 #

利用广播将x、y索引与偏移量结合，生成每个切片对应的所有行列索引：

# 生成所有行索引：形状为(len(x_indices), 2a+1)
x_grid = x_indices[:, np.newaxis] + offsets
# 生成所有列索引：形状为(len(x_indices), 2a+1)
y_grid = y_indices[:, np.newaxis] + offsets

这里通过[:, np.newaxis]给索引数组增加维度，让它能和offsets广播匹配。

步骤3：提取切片 #

由于需要每个切片是2D数组，需要将行列网格再扩展维度，让它们能正确对应每个切片的位置：

# 将x_grid扩展为(len(x_indices), 2a+1, 1)，y_grid扩展为(len(x_indices), 1, 2a+1)
# 广播后两者形状都是(len(x_indices), 2a+1, 2a+1)，对应每个切片的所有坐标
result = h[x_grid[:, :, np.newaxis], y_grid[:, np.newaxis, :]]

完整代码示例 #

import numpy as np

h = np.arange(49).reshape(7,7)
x_indices = np.array([1,3,4])
y_indices = np.array([2,3,5])
a = 1

offsets = np.arange(-a, a+1)
x_grid = x_indices[:, np.newaxis] + offsets
y_grid = y_indices[:, np.newaxis] + offsets

result = h[x_grid[:, :, np.newaxis], y_grid[:, np.newaxis, :]]
print(result.shape)  # (3, 3, 3)，符合预期
print(result)
# 输出：
# [[[ 1  2  3]
#   [ 8  9 10]
#   [15 16 17]]
# 
#  [[16 17 18]
#   [23 24 25]
#   [30 31 32]]
# 
#  [[25 26 27]
#   [32 33 34]
#   [39 40 41]]]

原理说明 #

• 广播机制让我们无需循环就能为每个中心坐标生成完整的切片索引网格
• 高级索引中，当两个索引数组的形状匹配时，会按对应位置提取元素，这里通过扩展维度让行列索引广播成每个切片的2D坐标网格，从而直接提取出每个正方形切片

内容的提问来源于stack exchange，提问作者OliverBunting