How can I use a 3d numpy array of indices to retrieve the corresponding values i a 4d array?

user_13045349 0
Viewed 86

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 Tis 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?

P.S. 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, e.g. the time derivative. Also, in reality the arrays are much larger then this minimal example.

numpy arrays python
2 Answers

Numpy advanced 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]
user_75708170

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)
user_17827920
Related