Numpy: for each element in one dimension, find coordinates of maximum of sub-array

Question

I've seen variations of this question asked a few times but so far haven't seen any answers that get to the heart of this general case. I have an n-dimensional array of shape [a, b, c, ...] . For some dimension x, I want to look at each sub-array and find the coordinates of the maximum.

For example, say b = 2, and that's the dimension I'm interested in. I want the coordinates of the maximum of [:, 0, :, ...] and [:, 1, :, ...] in the form a_max = [a_max_b0, a_max_b1], c_max = [c_max_b0, c_max_b1], etc.

I've tried to do this by reshaping my input matrix to a 2d array [b, a*c*d*...], using argmax along axis 0, and unraveling the indices, but the output coordinates don't wind up giving the maxima in my dataset. In this case, n = 3 and I'm interested in axis 1.

shape = gains_3d.shape
idx = gains_3d.reshape(shape[1], -1) 
idx = idx.argmax(axis = 1)
a1, a2 = np.unravel_index(idx, [shape[0], shape[2]])

Obviously I could use a loop, but that's not very pythonic.

For a concrete example, I randomly generated a 4x2x3 array. I'm interested in axis 1, so the output should be two arrays of length 2.

testarray = np.array([[[0.17028444, 0.38504759, 0.64852725],
        [0.8344524 , 0.54964746, 0.86628204]],

       [[0.77089997, 0.25876277, 0.45092835],
        [0.6119848 , 0.10096425, 0.627054  ]],

       [[0.8466859 , 0.82011746, 0.51123959],
        [0.26681694, 0.12952723, 0.94956865]],

       [[0.28123628, 0.30465068, 0.29498136],
        [0.6624998 , 0.42748154, 0.83362323]]])

testarray[:,0,:] is

array([[0.17028444, 0.38504759, 0.64852725],
       [0.77089997, 0.25876277, 0.45092835],
       [0.8466859 , 0.82011746, 0.51123959],
       [0.28123628, 0.30465068, 0.29498136]])

, so the first element of the first output array will be 2, and the first element of the other will be 0, pointing to 0.8466859. The second elements of the two matrices will be 2 and 2, pointing to 0.94956865 of testarray[:,1,:]

Please post a concrete MCVE with an example. I am having trouble understanding what you actually want. Some of the words say `gains_3d.argmax(1)`, but some don't — Mad Physicist, Jun 22 '21 at 20:11

hpaulj · Accepted Answer · 2021-06-22T21:07:14.553

Let's first try to get a clear idea of what you are trying to do:

Sample 3d array:

In [136]: arr = np.random.randint(0,10,(2,3,4))
In [137]: arr
Out[137]: 
array([[[1, 7, 6, 2],
        [1, 5, 7, 1],
        [2, 2, 5, *6*]],

       [[*9*, 1, 2, 9],
        [2, *9*, 3, 9],
        [0, 2, 0, 6]]])

After fiddling around a bit I came up with this iteration, showing the coordinates for each middle dimension, and the max value

In [151]: [(i,np.unravel_index(np.argmax(arr[:,i,:]),(2,4)),np.max(arr[:,i,:])) for i in range
     ...: (3)]
Out[151]: [(0, (1, 0), 9), (1, (1, 1), 9), (2, (0, 3), 6)]

I can move the unravel outside the iteration:

In [153]: np.unravel_index([np.argmax(arr[:,i,:]) for i in range(3)],(2,4))
Out[153]: (array([1, 1, 0]), array([0, 1, 3]))

Your reshape approach does avoid this loop:

In [154]: arr1 = arr.transpose(1,0,2)  # move our axis first
In [155]: arr1 = arr1.reshape(3,-1)    
In [156]: arr1
Out[156]: 
array([[1, 7, 6, 2, 9, 1, 2, 9],
       [1, 5, 7, 1, 2, 9, 3, 9],
       [2, 2, 5, 6, 0, 2, 0, 6]])

In [158]: np.argmax(arr1,axis=1)
Out[158]: array([4, 5, 3])
In [159]: np.unravel_index(_,(2,4))
Out[159]: (array([1, 1, 0]), array([0, 1, 3]))

max and argmax take only one axis value, where as you want the equivalent of taking the max along all but one axis. Some ufunc takes a axis tuple, but these do not. The transpose and reshape may be the only way.

In [163]: np.max(arr1,axis=1)
Out[163]: array([9, 9, 6])

Numpy: for each element in one dimension, find coordinates of maximum of sub-array

1 Answers1