如何使用 3d numpy 索引數組來檢索 4d 數組中的相應值？

Question

我有一個 4d numpy 數組temperature數據，在 x、y、z 點和時間 t 處測量溫度。 假設我有一個數組indices ，其中包含滿足條件的第一個實例的索引，例如temperature < 0 ，我如何提取滿足此條件的第一個溫度的 3d 數組？ 那就是我正在尋找相當於 numpy 的 1d 版本（默認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，則尺寸。

我當然可以用 for 循環來做到這一點：

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]]

但我正在尋找一些更優雅和更pythonic的東西。 有沒有更擅長 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)