Открыть сетки индексов массива

Question

Учитывая матрицу a где a.shape == (M, N, O), существует ли лучший способ генерации:

indices = (
    np.arange(M).reshape(M, 1, 1), 
    np.arange(N).reshape(1, N, 1), 
    np.arange(O).reshape(1, 1, O) 
) 

А также работает для более высоких измерений массива?

Я могу приблизиться к np.indices, но это возвращает замкнутую сетку (все записи имеют форму M,N,O), а не открытую сетку.

Answer 1

Возможно, вы можете сделать это с np.ix_:

np.ix_(np.arange(M), np.arange(N), np.arange(O)) 

Из документации:

Построить открытую сетку из нескольких последовательностей.

Эта функция принимает N 1-D последовательности и возвращает N выходов с N измерениями каждый, таким образом, что форма равна 1 во всех измерениях, кроме одного измерения, а размер с величиной без единицы измеряется во всех N измерениях.

Это дает тот же результат, как indices здесь (M, N, O = 2, 3, 4):

(array([[[0]], 

     [[1]]]), 

array([[[0], 
     [1], 
     [2]]]), 

array([[[0, 1, 2, 3]]])) 

Answer 2

np.ogrid является, вероятно, самым кратким способом, если M, N и O указаны отдельно:

M, N, O = 2, 3, 4 

indices = np.ogrid[:M, :N, :O] 

print(indices) 
# [array([[[0]], 

#  [[1]]]), array([[[0], 
#   [1], 
#   [2]]]), array([[[0, 1, 2, 3]]])] 

Если вход - это только массив, самый сжатый способ, который я могу придумать:

np.ix_(*(np.r_[:s] for s in A.shape))