Finding index of maximum value in array with NumPy

Question

I would like to find a maximum in a float64 array, excluding nan values.

I saw np.nanmax function but it doesn't give the index corresponding to the found value.

it 's quite strange to scan after to the value specially the function necessarily use the index ??? Can't it be a mistake searching like that .

isn't there a way to recover the index directly ?

Answer 1

Numpy has an argmax function that returns just that, although you will have to deal with the nans manually. nans always get sorted to the end of an array, so with that in mind you can do:

a = np.random.rand(10000)
a[np.random.randint(10000, size=(10,))] = np.nan
a = a.reshape(100, 100)

def nanargmax(a):
    idx = np.argmax(a, axis=None)
    multi_idx = np.unravel_index(idx, a.shape)
    if np.isnan(a[multi_idx]):
        nan_count = np.sum(np.isnan(a))
        # In numpy < 1.8 use idx = np.argsort(a, axis=None)[-nan_count-1]
        idx = np.argpartition(a, -nan_count-1, axis=None)[-nan_count-1]
        multi_idx = np.unravel_index(idx, a.shape)
    return multi_idx

>>> nanargmax(a)
(20, 93)

Answer 2

You should use np.where

In [17]: a=np.random.uniform(0, 10, size=10)

In [18]: a
Out[18]: 
array([ 1.43249468,  4.93950873,  7.22094395,  1.20248629,  4.66783985,
        6.17578054,  4.6542771 ,  7.09244492,  7.58580515,  5.72501954])

In [20]: np.where(a==a.max())
Out[20]: (array([8]),)

This also works for 2 arrays, the returned value, is the index. Here we create a range from 1 to 9:

 x = np.arange(9.).reshape(3, 3)

This returns the index, of the the items that equal 5:

In [34]: np.where(x == 5)
Out[34]: (array([1]), array([2])) # the first one is the row index, the second is the column

You can use this value directly to slice your array:

In [35]: x[np.where(x == 5)]
Out[35]: array([ 5.])