如何对未知维度的ndarray进行动态范围索引？

问题背景 #

我有形状未知的data数组，以及用于对data进行切片的bounds数组。现有代码仅适配3D的data，需要将其泛化为支持任意N维数组。尝试使用Python的range对象作为索引，但未得到预期效果。

示例数据 #

import numpy as np

# 2D示例数组
data2D = np.arange(2*3).reshape((2, 3))
# 3D示例数组
data3D = np.arange(2*3*4).reshape((2, 3, 4))

# 2D数组对应的切片边界
bounds2D = np.array([[[0, 2], [0, 2]], [[0, 2], [1, 3]]])
# 3D数组对应的切片边界
bounds3D = np.array(
    [
        [[0, 2], [0, 2], [0, 2]],
        [[0, 2], [0, 2], [2, 4]],
        [[0, 2], [1, 3], [0, 2]],
        [[0, 2], [1, 3], [2, 4]],
    ],
)

原3D处理代码 #

for b in bounds:
    l0, u0 = b[0]
    l1, u1 = b[1]
    l2, u2 = b[2]
    a = data[l0:u0, l1:u1, l2:u2]
    print(a)

解决方案 #

核心是动态生成切片元组——numpy数组的索引接受由slice对象组成的元组，我们只需遍历每个维度的边界，将[lower, upper]转换为slice(lower, upper)，再打包成元组即可。

实现代码：

import numpy as np

def slice_n_dim_data(data, bounds):
    for b in bounds:
        # 为每个维度生成slice对象，打包成元组
        slice_tuple = tuple(slice(lower, upper) for lower, upper in b)
        sliced_result = data[slice_tuple]
        print(sliced_result)

# 测试2D场景
slice_n_dim_data(data2D, bounds2D)
print("---")
# 测试3D场景
slice_n_dim_data(data3D, bounds3D)

关于range失效的原因 #

numpy不支持用range对象做维度切片：range会被识别为花式索引（按序列中的整数逐个取元素），而非连续区间的切片操作。必须使用slice对象（也就是代码中a:b语法对应的底层对象）才能实现连续切片。

内容的提问来源于stack exchange，提问作者Nikola Jovišić