部分索引多维数组

Question

import numpy as np
arr = np.random.rand(50,3,3,3,16)
ids = (0,0,2,10)
b = arr[:, ids]  # don't work
b = arr[:, *ids]  # don't work
b = arr[:][ids]  # don't work
b = arr[:, tuple(ids)]  # don't work
b = arr[: + ids]  # don't work, obviously..
# b = arr[:,0,0,2,10].shape  # works (desired outcome)

我知道对此有几个问题，例如元组作为多维数组的索引，或者拆包元组/数组/列表作为Numpy数组的索引，但是它们都不适合我的情况。 基本上，我想在第一个轴上索引所有其余轴中指定的“列”中的所有内容（请参阅代码的最后一行）。 在这种情况下(50,)所需的输出形状应为(50,) 。

但是我想用元组/ ID列表建立索引，因为我需要遍历它们，例如：

all_ids = ((0,0,0,2), (0,0,0,6), (1,1,0,2), (1,1,0,6),
           (2,2,0,2), (2,2,0,6), (2,2,2,2), (2,2,2,6))
c = 0
for id in all_ids:
    c += arr[:, id].sum() 

Answer 1

将slice(None)添加到ids第一维，然后添加子集：

arr[(slice(None),) + ids].shape
# (50,)

哪里：

(slice(None),) + ids
# (slice(None, None, None), 0, 0, 2, 10)

注意slice(None, None, None)等效于: ，即全部切片。 您可以在此处阅读有关使用slice对象进行索引的文档 。