Index of maximum values along and plane in a numpy 3D array

Question

I have a 3D numpy array of shape (3,3,3). I would like to obtain indices of maximum values in a plane,"plane" according to me is as follows:

a = np.random.rand(3,3,3)
>>> a[:,:,0]
array([[0.98423332, 0.44410844, 0.06945133],
       [0.69876575, 0.87411547, 0.53595041],
       [0.53418486, 0.16186808, 0.60579623]])
>>> a[:,:,1]
array([[0.38969199, 0.80202126, 0.62189662],
       [0.66609605, 0.09771614, 0.74061269],
       [0.77081531, 0.20068743, 0.72762023]])
>>> a[:,:,2]
array([[0.57110332, 0.29021439, 0.15433043],
       [0.21762439, 0.93112448, 0.05763075],
       [0.77880124, 0.36637245, 0.29070822]])

I have a solution but I would like to have something shorter and quicker without for loops, my solution is as belows:

for i in range(3):
    x=a[:,:,i].argmax()/3
    y=a[:,:,i].argmax()%3
    z=i
    print(x,y,z)
    print a[x][y][z]
(0, 0, 0)
0.9842333247061394
(0, 1, 1)
0.8020212566990867
(1, 1, 2)
0.9311244845473187

np.argmax() has an axis keyword argument you should look into — Primusa
– Primusa, Commented Jul 19, 2018 at 15:30
I don't this that will help, it will give a 3x3 array of max elements for all the axis cases i.e 0,1,2. — pythonic
– pythonic, Commented Jul 19, 2018 at 16:24

Divakar · Accepted Answer · 2018-07-19 15:58:33Z

2

We simply need to reshape the input array to 2D by merging the last two axes and then applying argmax along the second one i.e. the merged one to give ourselves a vectorized approach -

def argmax_each_plane(a):
    a2D = a.reshape(a.shape[0],-1)
    idx = a2D.argmax(1)
    indices = np.unravel_index(idx, a.shape[1:])
    vals = a2D[np.arange(len(idx)), idx]
    return vals, np.c_[indices]

Sample run -

In [60]: np.random.seed(0)
    ...: a = np.random.rand(3,3,3)

In [61]: a
Out[61]: 
array([[[0.5488135 , 0.71518937, 0.60276338],
        [0.54488318, 0.4236548 , 0.64589411],
        [0.43758721, 0.891773  , 0.96366276]],

       [[0.38344152, 0.79172504, 0.52889492],
        [0.56804456, 0.92559664, 0.07103606],
        [0.0871293 , 0.0202184 , 0.83261985]],

       [[0.77815675, 0.87001215, 0.97861834],
        [0.79915856, 0.46147936, 0.78052918],
        [0.11827443, 0.63992102, 0.14335329]]])

In [62]: v, ind = argmax_each_plane(a)

In [63]: v
Out[63]: array([0.96366276, 0.92559664, 0.97861834])

In [64]: ind
Out[64]: 
array([[2, 2],
       [1, 1],
       [0, 2]])

If you need z indices as well, use : np.c_[indices[0], indices[1], range(len(a2D))].

edited Jul 19, 2018 at 15:58

answered Jul 19, 2018 at 15:40

Divakar

222k19 gold badges273 silver badges374 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

pythonic Over a year ago

this didn't work, i want maximum of a[:,:,0] , a[:,:,1], a[:,:,2], note that my "planes" are different from the "planes" that numpy prints when you print a.

Divakar Over a year ago

@pythonic Then simply use : argmax_each_plane(np.moveaxis(a,-1,0)).

Collectives™ on Stack Overflow

Index of maximum values along and plane in a numpy 3D array

1 Answer 1

2 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

2 Comments

Your Answer

Sign up or log in

Post as a guest

Related