将每个跨步滚动窗口的第一个元素引用回源数组的更好方法？

Question

我具有下面定义的滚动窗口功能，以“滚动”通过一个numpy数组。功能定义如下。

roller = np.arange(1,21).reshape(4,5)

def rolling_window(main_array, sub_array_shape, stepsize_x=1, stepsize_y=1): 
    strided = np.lib.stride_tricks.as_strided
    x_sub_dim, y_sub_dim = sub_array_shape  #Define custom sub_array_shape for rolling window.  
    x_main_dim,y_main_dim = main_array.shape[-2:]    # List slice ensures that multi-dimension is not affected
    x_main_stride,y_main_stride = main_array.strides[-2:]

    out_shp = main_array.shape[:-2] + (x_main_dim - x_sub_dim + 1, y_main_dim - y_sub_dim + 1, x_sub_dim, y_sub_dim)
    out_stride = main_array.strides[:-2] + (x_main_stride, y_main_stride, x_main_stride, y_main_stride)

    imgs = strided(main_array, shape=out_shp, strides=out_stride)
    return imgs[...,::stepsize_x,::stepsize_y,:,:]

rolling_window(roller, (2,3))



Output: 
[[[[ 1  2  3]
   [ 6  7  8]]

  [[ 2  3  4]
   [ 7  8  9]]

  [[ 3  4  5]
   [ 8  9 10]]]


 [[[ 6  7  8]
   [11 12 13]]

  [[ 7  8  9]
   [12 13 14]]

  [[ 8  9 10]
   [13 14 15]]]


 [[[11 12 13]
   [16 17 18]]

  [[12 13 14]
   [17 18 19]]

  [[13 14 15]
   [18 19 20]]]]

我想获取每个子数组的第一个值（例如此输出），可以通过将return slice更改为

来实现。

return imgs[...,:stepsize_x,:stepsize_y]
Output:
[[[[ 1]]

  [[ 2]]

  [[ 3]]]


 [[[ 6]]

  [[ 7]]

  [[ 8]]]


 [[[11]]

  [[12]]

  [[13]]]]

我需要从子数组中引用所有这些第一元素回到主数组的位置。我正在考虑将输出值附加到列表中，然后使另一个函数搜索主数组以查找每个第一个数字的位置。是在一个函数中执行此操作的更优雅的方法，还是正确的调用将其拆分为多个函数？