Make np.argmax, np.argmin to return indexes and values at the same time

Actually, np.argmax returns indexes of maximum values along some axis. On the other side, np.amax returns maximum values along some axis but if you want to get indexes of maximum values with those maximum values it is necessary to go through the array again. For example:

import numpy as np

data = np.array([[0.02079073, 0.97920927], [0.00487725, 0.99512275], [0.00849596, 0.99150404], [0.96402552, 0.03597448], [0.00711506, 0.99288494]])

indexes = np.argmax(data, axis=-1)

maximum_values = data[np.arange(data.shape[0]), indexes]

print(indexes)
print(maximum_values)
print(np.amax(data, axis=-1))
>> [1 1 1 0 1]
>> [0.97920927 0.99512275 0.99150404 0.96402552 0.99288494]
>> [0.97920927 0.99512275 0.99150404 0.96402552 0.99288494]

Even with np.arange, go through the array again consumes time (specially on huge arrays). So, np.argmax and np.argmin could return indexes and values at the same time on one look (maybe with an extra parameter in order to make compatible previous versions). For 1D just indexing would be fine, but that not work for arrays with 2 or more axis because it gets the rows elements:

import numpy as np

data = np.array([0.3, 0.7, 0.5])

indexes = np.argmax(data, axis=-1)

print(data[indexes])
>> 0.7

import numpy as np
data = np.array([[0.02079073, 0.97920927], [0.00487725, 0.99512275], [0.00849596, 0.99150404], [0.96402552, 0.03597448], [0.00711506, 0.99288494]])

indexes = np.argmax(data, axis=-1)

print(data[indexes])
>> [[0.00487725 0.99512275]
      [0.00487725 0.99512275]
      [0.00487725 0.99512275]
      [0.02079073 0.97920927]
      [0.00487725 0.99512275]]

Numpy version: ‘1.17.5’ Python version: ‘3.6.9’

[WarrenWeckesser]

If you really want to do this in a single function call, over in ufunclab I have implemented an assortment of NumPy ufuncs and gufuncs, including min_argmin and max_argmax.

Here’s your example, using ufunclab.max_argmax:

In [21]: import numpy as np                                                     

In [22]: import ufunclab                                                        

In [23]: data = np.array([[0.02079073, 0.97920927], [0.00487725, 0.99512275], [0
    ...: .00849596, 0.99150404], [0.96402552, 0.03597448], [0.00711506, 0.992884
    ...: 94]])                                                                  

In [24]: data                                                                   
Out[24]: 
array([[0.02079073, 0.97920927],
       [0.00487725, 0.99512275],
       [0.00849596, 0.99150404],
       [0.96402552, 0.03597448],
       [0.00711506, 0.99288494]])

In [25]: values, indices = ufunclab.max_argmax(data, axis=-1)                   

In [26]: values                                                                 
Out[26]: array([0.97920927, 0.99512275, 0.99150404, 0.96402552, 0.99288494])

In [27]: indices                                                                
Out[27]: array([1, 1, 1, 0, 1])

[gmoda]

You could also do:

[ind1[ind2] for ind1, ind2 in zip(data, indexes)]