Is there a Julia analogue to numpy.argmax?

Question

In Python, there is numpy.argmax:

In [7]: a = np.random.rand(5,3)

In [8]: a
Out[8]: 
array([[ 0.00108039,  0.16885304,  0.18129883],
       [ 0.42661574,  0.78217538,  0.43942868],
       [ 0.34321459,  0.53835544,  0.72364813],
       [ 0.97914267,  0.40773394,  0.36358753],
       [ 0.59639274,  0.67640815,  0.28126232]])

In [10]: np.argmax(a,axis=1)
Out[10]: array([2, 1, 2, 0, 1])

Is there a Julia analogue to Numpy's argmax? I only found a indmax, which only accept a vector, not a two dimensional array as np.argmax.

[There is one now, in 0.7-alpha](https://stackoverflow.com/a/50955524/3565696) — ederag, Jun 20 '18 at 19:18

tholy · Accepted Answer · 2020-03-03T19:24:07.260

The fastest implementation will usually be findmax (which allows you to reduce over multiple dimensions at once, if you wish):

julia> a = rand(5, 3)
5×3 Array{Float64,2}:
 0.867952  0.815068   0.324292
 0.44118   0.977383   0.564194
 0.63132   0.0351254  0.444277
 0.597816  0.555836   0.32167 
 0.468644  0.336954   0.893425

julia> mxval, mxindx = findmax(a; dims=2)
([0.8679518267243425; 0.9773828942695064; … ; 0.5978162823947759; 0.8934254589671011], CartesianIndex{2}[CartesianIndex(1, 1); CartesianIndex(2, 2); … ; CartesianIndex(4, 1); CartesianIndex(5, 3)])

julia> mxindx
5×1 Array{CartesianIndex{2},2}:
 CartesianIndex(1, 1)
 CartesianIndex(2, 2)
 CartesianIndex(3, 1)
 CartesianIndex(4, 1)
 CartesianIndex(5, 3)

`ind2sub` has been removed in 0.7: https://discourse.julialang.org/t/psa-replacement-of-ind2sub-sub2ind-in-julia-0-7/14666, use CartesianIndex instead. — BoZenKhaa, Mar 02 '20 at 17:53

score 3 · Answer 2 · edited Jun 20 '20 at 09:12

According to the Numpy documentation, argmax provides the following functionality:

numpy.argmax(a, axis=None, out=None)

Returns the indices of the maximum values along an axis.

I doubt a single Julia function does that, but combining mapslices and argmax is just the ticket:

julia> a = [ 0.00108039  0.16885304  0.18129883;
             0.42661574  0.78217538  0.43942868;
             0.34321459  0.53835544  0.72364813;
             0.97914267  0.40773394  0.36358753;
             0.59639274  0.67640815  0.28126232] :: Array{Float64,2}

julia> mapslices(argmax,a,dims=2)
5x1 Array{Int64,2}:
 3
 2
 3
 1
 2

Of course, because Julia's array indexing is 1-based (whereas Numpy's array indexing is 0-based), each element of the resulting Julia array is offset by 1 compared to the corresponding element in the resulting Numpy array. You may or may not want to adjust that.

If you want to get a vector rather than a 2D array, you can simply tack [:] at the end of the expression:

julia> b = mapslices(argmax,a,dims=2)[:]
5-element Array{Int64,1}:
 3
 2
 3
 1
 2

Since Julia 0.7, use `argmax` instead of `indmax`: https://stackoverflow.com/questions/37474493/location-of-minimum-in-julia — BoZenKhaa, Mar 02 '20 at 19:23

score 1 · Answer 3 · answered Mar 03 '20 at 14:10

To add to the jub0bs's answer, argmax in Julia 1+ mirrors the behavior of np.argmax, by replacing axis with dims keyword, returning CarthesianIndex instead of index along given dimension:

julia>  a = [ 0.00108039  0.16885304  0.18129883;

                0.42661574  0.78217538  0.43942868;      

                0.34321459  0.53835544  0.72364813;      

                0.97914267  0.40773394  0.36358753;      

                0.59639274  0.67640815  0.28126232] :: Array{Float64,2}

julia> argmax(a, dims=2)
5×1 Array{CartesianIndex{2},2}:
CartesianIndex(1, 3)
CartesianIndex(2, 2)
CartesianIndex(3, 3)
CartesianIndex(4, 1)
CartesianIndex(5, 2)

Is there a Julia analogue to numpy.argmax?

3 Answers3