Как я могу использовать массив индексов 3d numpy для получения соответствующих значений в массиве 4d?

Question

У меня есть массив 4d numpy temperature данных с измеренной температурой в точках x, y, z и времени t. Предполагая, что у меня есть массив indices с индексами, где встречается первый экземпляр условия, скажем, temperature < 0, как мне извлечь трехмерный массив с первыми температурами, удовлетворяющими этому условию? То есть я ищу эквивалент 1d версии numpy (import numpy as np молчаливо предполагается)

>>> temperatures = np.arange(10,-10,-1)
>>> ind = np.argmax(temperatures < 0)
>>> T = temperature[ind]

Я пробовал аналогичный

In [1]: temperatures = np.random.random((11,8,5,200)) * 1000

In [2]: temperatures.shape
Out[2]: (11, 8, 5, 200)

In [3]: indices= np.argmax(temperatures > 900,axis=3)

In [4]: indices.shape
Out[4]: (11, 8, 5)

In [5]: T = temperatures[:,:,:,indices]

In [6]: T.shape
Out[6]: (11, 8, 5, 11, 8, 5)

Однако, размеры, если T равно 6.

Конечно, я мог бы сделать это с помощью l oop:

indices = np.argmax(temperatures > 900,axis=3)
x,y,z = temperatures.shape[:-1]
T = np.zeros((x,y,z))
for indx in range(x):
    for indy in range(y):
        for indz in range(z):
            T[indx,indy,indz] = temperatures[indx,indy,indz,indices[indx,indy,indz]]

, но я ищу что-то элегантное и более питоническое c. Есть ли кто-то более опытный с numpy, который может мне помочь в этом?

PS Ради ясности, я не просто ищу температуру в этих точках, заданную indices, Я также ищу другие величины в массивах той же формы, что и temperature, например производная по времени. Кроме того, в действительности массивы намного больше, чем этот минимальный пример.

Answer 1

Numpy расширенное индексирование всегда работает:

import numpy as np 
temperatures = np.random.random((11,8,5, 200)) * 1000
indices = np.argmax(temperatures > 900, axis=3)

x, y, z = temperatures.shape[:-1]

T = temperatures[np.arange(x)[:, np.newaxis, np.newaxis],
                 np.arange(y)[np.newaxis, :, np.newaxis],
                 np.arange(z)[np.newaxis, np.newaxis, :],
                 indices]

Как указала jdehesa, это можно сделать более кратким:

x, y, z = np.ogrid[:x, :y, :z]
T = temperatures[x, y, z, i]

Answer 2

Я думаю, вам нужно:

axis = 3
indices = np.argmax(temperatures > 900, axis=axis)
result = np.take_along_axis(temperatures, np.expand_dims(indices, axis), axis)
result = result.squeeze(axis)