Search code examples
pythonarraysnumpymatrixargmax

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


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


Solution

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