从维数未知的 numpy 数组中提取超立方块

Question

我有一些 Python 代码，目前与二维数组硬连线如下：

import numpy as np
data = np.random.rand(5, 5)
width = 3

for y in range(0, data.shape[1] - W + 1):
    for x in range(0, data.shape[0] - W + 1):
        block = data[x:x+W, y:y+W]
        # Do something with this block

现在，这是为二维数组硬编码的，我想将其扩展到 3D 和 4D 数组。当然，我可以为其他维度编写更多函数，但我想知道是否有 python/numpy 技巧来生成这些子块，而不必为多维数据复制此函数。

Answer 1

这是我对这个问题的看法。下面代码背后的想法是找到每个数据片的“起始索引”。因此，对于 5x5x5 数组的 4x4x4 子数组，起始索引为 (0,0,0), (0,0,1), (0,1,0), (0,1,1), (1,0,0), (1,0,1), (1,1,1)，沿每个维度的切片长度为 4。

要获取子数组，您只需要遍历切片对象的不同元组并将它们传递给数组。

import numpy as np
from itertools import product

def iterslice(data_shape, width):
    # check for invalid width
    assert(all(sh>=width for sh in data_shape), 
           'all axes lengths must be at least equal to width')

    # gather all allowed starting indices for the data shape
    start_indices = [range(sh-width+1) for sh in data_shape]

    # create tuples of all allowed starting indices
    start_coords = product(*start_indices)

    # iterate over tuples of slice objects that have the same dimension
    # as data_shape, to be passed to the vector
    for start_coord in start_coords:
        yield tuple(slice(coord, coord+width) for coord in start_coord)

# create 5x5x5 array
arr = np.arange(0,5**3).reshape(5,5,5)

# create the data slice tuple iterator for 3x3x3 sub-arrays
data_slices = iterslice(arr.shape, 3)

# the sub-arrays are a list of 3x3x3 arrays, in this case
sub_arrays = [arr[ds] for ds in data_slices]