2

I have a 3-d array. I find the indexes of the maxima along an axis using argmax. How do I now use these indexes to obtain the maximal values? 2nd part: How to do this for arrays of N-d?

Eg:

u = np.arange(12).reshape(3,4,1)

In [125]: e = u.argmax(axis=2)
Out[130]: e
array([[0, 0, 0, 0],
   [0, 0, 0, 0],
   [0, 0, 0, 0]])

It would be nice if u[e] produced the expected results, but it doesn't work.

nbecker
  • 1,645
  • 5
  • 17
  • 23
  • Doing `argmax(2)` makes no sense since the 3'rd dimension only has one element. – cs95 Mar 20 '18 at 13:19
  • `u[e]` will give you the maximum values in third axis but since each row at this axis has only one element in it, it returns the index of that only element which is 0. – Mazdak Mar 20 '18 at 13:53
  • I maybe confused people with the comment "It would be nice...". That's not the question. My question is "How do I now use the indexes returned by argmax to get the maximal values?" And, "How to generalize to arrays of N-d?" – nbecker Mar 20 '18 at 15:47

2 Answers2

2

The return value of argmax along an axis can't be simply used as an index. It only works in a 1d case.

In [124]: u = np.arange(12).reshape(3,4,1)
In [125]: e = u.argmax(axis=2)
In [126]: u.shape
Out[126]: (3, 4, 1)
In [127]: e.shape
Out[127]: (3, 4)

e is (3,4), but its values only index the last dimension of u. 

In [128]: u[e].shape
Out[128]: (3, 4, 4, 1)

Instead we have to construct indices for the other 2 dimensions, ones which broadcast with e. For example:

In [129]: I,J=np.ix_(range(3),range(4))
In [130]: I
Out[130]: 
array([[0],
       [1],
       [2]])
In [131]: J
Out[131]: array([[0, 1, 2, 3]])

Those are (3,1) and (1,4). Those are compatible with (3,4) e and the desired output

In [132]: u[I,J,e]
Out[132]: 
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])

This kind of question has been asked before, so probably should be marked as a duplicate. The fact that your last dimension is size 1, and hence e is all 0s, distracting readers from the underlying issue (using a multidimensional argmax as index).

numpy: how to get a max from an argmax result

Get indices of numpy.argmax elements over an axis

Assuming you've taken the argmax on the last dimension

In [156]: ij = np.indices(u.shape[:-1])
In [157]: u[(*ij,e)]
Out[157]: 
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])

or:

ij = np.ix_(*[range(i) for i in u.shape[:-1]])

If the axis is in the middle, it'll take a bit more tuple fiddling to arrange the ij elements and e.

hpaulj
  • 221,503
  • 14
  • 230
  • 353
  • Same principle applies - generate broadcastable indices to work with `e`. `*` unpacking makes this easier in some cases. See my edits. – hpaulj Mar 20 '18 at 18:47
0

so for general N-d array

dims = np.ix_(*[range(x) for x in u.shape[:-1]])
u.__getitem__((*dims,e))

You can't write u[*dims,e], that's a syntax error, so I think you must use getitem directly.

nbecker
  • 1,645
  • 5
  • 17
  • 23
  • `u[(*ij,e)]` works. `u[i,j]` is the same as `u[(i,j)]`, that is there's an implicit tuple. So making it explicit is fine. – hpaulj Mar 20 '18 at 18:52