python - 使用变量 NumPy 切片保留维度

Question

我有一个 3d NumPy 数组，我想从中获取许多切片。这些切片在第一维和第二维中的长度将为一个或多个，而第三个维度将完整返回。切片应始终为 3d。

我的尝试:

import numpy as np

a = np.zeros((1000, 10, 100))
row_sets = ([19, 20], [21])
col_sets = ([6], [7, 8])

for rows in row_sets:
    for cols in col_sets:
        b = a[[rows], [cols]]
        print(rows, cols, b.shape)

结果:

[19, 20] [6] (1, 2, 100)
[19, 20] [7, 8] (1, 2, 100)
[21] [6] (1, 1, 100)
[21] [7, 8] (1, 2, 100)

如果我从切片中删除嵌套括号:

b = a[rows, cols]

我在第二个维度中遇到了同样的问题，并且维度没有保留:

[19, 20] [6] (2, 100)
[19, 20] [7, 8] (2, 100)
[21] [6] (1, 100)
[21] [7, 8] (2, 100)

我正在寻找的结果是这样的:

[19, 20] [6] (2, 1, 100)
[19, 20] [7, 8] (2, 2, 100)
[21] [6] (1, 1, 100)
[21] [7, 8] (1, 2, 100)

Answer 1

您正在触发 advanced indexing 通过使用整数列表作为索引，这会减少结果数组的维度，如果您仍想对数组进行切片，可以使用 np.ix_从整数列表重建切片索引:

for rows in row_sets:
    for cols in col_sets:
        b = a[np.ix_(rows, cols)]
        print(rows, cols, b.shape)

#[19, 20] [6] (2, 1, 100)
#[19, 20] [7, 8] (2, 2, 100)
#[21] [6] (1, 1, 100)
#[21] [7, 8] (1, 2, 100)