python - Getting indices from numpy array and applying to 2nd lower dimensional array to create new array -

i have array gamma, gamma.shape=(20,7,90,144), represents dimensions (t,z,y,x) , respective indices denote [l,k,j,i]. each (t,y,x), want find lowest value of k such gamma[l,k,j,i] > g_crit, g_crit constant. denote set of lowest k values k_low; i've found can nicely

k_low = np.argmax(gamma > g_crit, axis=1) 

there array levs, levs.shape=(7,), corresponding z dimension well. trying create array levs_low, such levs_low.shape=(20,90,144) , levs_low[l,j,i]=levs[k_low[l,j,i]]. i'm stuck on step. ideas please? thanks.

this should trick:

levs_low=levs[k_low]  >>> gamma=np.random.rand(20,7,90,144) >>> k_low = np.argmax(gamma > .3, axis=1) >>> levs=np.random.rand(7) >>> levs_low=levs[k_low] >>> levs_low.shape (20, 90, 144) 

for small example:

>>> g=np.random.randint(0,5,(4,4)) >>> g array([[2, 0, 2, 2],        [2, 0, 1, 0],        [3, 3, 0, 3],        [3, 0, 4, 4]]) >>> k=np.arange(5)*-1 >>> k array([ 0, -1, -2, -3, -4]) >>> k[g]   #uses indices of g select values k. same np.take(k,g) array([[-2,  0, -2, -2],        [-2,  0, -1,  0],        [-3, -3,  0, -3],        [-3,  0, -4, -4]]) 

@saullo castro's answer interesting. strange there ~5x difference in timings between fancy indexing , np.take.

%timeit levs[k_low] 100 loops, best of 3: 2.3 ms per loop  %timeit np.take( levs, k_low ) 1000 loops, best of 3: 439 per loop  in [33]: np.all(levs[k_low]==np.take(levs,k_low)) out[33]: true