Частичное индексирование многомерного массива

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 Arrays , но ни один из них не работает для мое дело. По сути, я хочу индексировать все по первой оси, из указанных «столбцов» в остальных осях (см. Последнюю строку моего кода). В этом случае желаемая выходная форма должна быть (50,).

Но я хочу индексировать с помощью кортежа / списка идентификаторов, потому что мне нужно перебирать их, например:

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 объекта для индексации здесь .