python - Numpy 同时切片多个维度

Question

给定一个二维数组和两对索引，分别定义子矩阵的左上角和右下角:

a = np.arange(25).reshape(5,5)
# array([[ 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]])
ij_ul = (1, 1)  # upper left
ij_lr = (4, 3)  # lower right

我可以通过以下方式获取子矩阵:

a[ij_ul[0]:ij_lr[0], ij_ul [1]: ij_lr[1]]
# array([[ 6,  7],
#        [11, 12],
#        [16, 17]])

但是我想使用像这样更优雅的东西:

a[ij_ul:ij_lr]  # Wish
# TypeError: slice indices must be integers or None or have an __index__ method

同时切片所有维度。 ij_ul 中的第一个索引切片到 ij_lr 中的第一个索引，第二个索引切片到第二个索引。

我想将其用于应适用于 2D 和 3D 的函数。这里这样的公式会更通用，因为不需要计算维度和调整切片:

a[i:j]  # 1D
a[i:j, k:l]  # 2D
a[i:j, k:l, n:m]  # 3D
a[tuple_1:tuple_2]  # Wish

如果数组和两个索引元组是，最后一行总是有效匹配尺寸。

在 python/numpy 中是否有类似同时切片多个维度的东西？

Answer 1

鉴于您有这两个列表，我们可以构造一个元组，如下所示:

our_submatrix = a[tuple(map(slice, ij_ul, ij_lr))]

在这里，我们同时迭代两个可迭代对象 ij_ul 和 ij_lr，并构造 slice(..) 对象(如果您编写i:j 你实际上在幕后构造了这样的对象)。

所以我们在这里做的是构造一个切片元组，我们可以将其传递给 a 来构造具有指定切片的 View 。

例如:

>>> A = np.random.randn(5, 7)
>>> A
array([[-0.41526531,  0.88895983,  1.25061578, -0.01930474,  0.90183725,
         0.169831  , -0.67271614],
       [ 0.26148559, -1.24269541,  0.72472787, -2.75687495, -0.64511098,
        -1.5535446 ,  0.74714421],
       [ 0.8423359 ,  0.26620621,  0.26581058, -1.20667371,  0.60521744,
         2.02856815, -0.43339106],
       [ 0.29617257,  1.86195973, -1.33591589,  0.97233933, -0.69254736,
        -0.31652783,  0.78613453],
       [-0.93928881,  0.26735583, -0.27057705,  0.78084479, -0.76763169,
         0.12574734, -1.88236193]])
>>> A[tuple(map(slice, ij_ul, ij_lr))]
array([[-1.24269541,  0.72472787],
       [ 0.26620621,  0.26581058],
       [ 1.86195973, -1.33591589]])