如何使用 3d numpy 索引数组来检索 4d 数组中的相应值？

Question

I have a 4d numpy array temperature of data with the measured temperature at points x,y,z and time t. Assuming I have an array indices with the indices where the first instance of a condition is met, say temperature < 0 , how do I extract a 3d array with the first temperatures satisfying this condition? That is I'm looking for the equivalent of numpy's 1d version ( import numpy as np tacitly assumed)

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

I have tried the analogous

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)

However, the dimensions if T is 6.

I could of course do it with a for loop:

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

but I'm looking for something fore elegant and more pythonic. Is there someone more skilled with numpy out there who can help me out on this?

PS For the sake of clarity, I'm not just looking for the temperature at these points given by indices , I'm also looking for other quantities in arrays of the same shape as temperature , eg the time derivative. Also, in reality the arrays are much larger then this minimal example.

Answer 1

Numpyadvanced indexing does always work:

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]

As jdehesa pointed out this can be made more concise:

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

Answer 2

I think you need:

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)