[Numpy-discussion] Extracting values from one array corresponding to argmax elements in another array

[josef.pktd at gmail.com]

On Tue, Apr 6, 2010 at 9:22 AM, Ken Basye <kbasye1 at jhu.edu> wrote:
> From: Vincent Schut <schut at sarvision.nl>
>
> On 04/05/2010 06:06 PM, Keith Goodman wrote:
>
>
> On Mon, Apr 5, 2010 at 8:44 AM, Ken Basye<kbasye1 at jhu.edu>  wrote:
>
>
> Hi Folks,
>   I have two arrays, A and B, with the same shape.  I want to find the
> highest values in A along some axis, then extract the corresponding
> values from B.  I can get the highest values in A with A.max(axis=0) and
> the indices of these highest values with A.argmax(axis=0).  I'm trying
> to figure out a loop-free way to extract the corresponding elements from
> B using these indices.  Here's code with a loop that will do what I want
> for two-dimensional arrays:
>
>   >>>  a
> array([[ 100.,    0.,    0.],
>        [   0.,  100.,  100.],
>        [   0.,    0.,    0.]])
>
>   >>>  a.max(axis=0)
> array([ 100.,  100.,  100.])
>
>   >>>  sel = a.argmax(axis=0)
>   >>>sel
> array([0, 1, 1])
>
>   >>>  b = np.arange(9).reshape((3,3))
>   >>>  b
> array([[0, 1, 2],
>        [3, 4, 5],
>        [6, 7, 8]])
>
>   >>>  b_best = np.empty(3)
>   >>>  for i in xrange(3):
> ...    b_best[i] = b[sel[i], i]
> ...
>   >>>  b_best
> array([ 0.,  4.,  5.])
>
>
> Here's one way:
>
>
>
> b[a.argmax(axis=0), range(3)]
>
>
>     array([0, 4, 5])
>
>
> Which does not work anymore when your arrays become more-dimensional
> (like in my case: 4 or more) and the axis you want to select on is not
> the first/last one. If I recall correctly, I needed to construct the
> full index arrays for the other dimensions too (like with ogrid I
> think). So: create the ogrid, replace the one for the dimensions you
> want the argmax selection to take place on with the argmax parameter,
> and use those index arrays to index your b array.
> I'd need to look up my source code to be more sure/precise. If anyone
> would like me to, please let me know. If anyone knows a less elaborate
> way, also please let us know! :-)
>
>
> Hi Vincent,
>   I'd like to see more about your solution.  For my present purposes,
> Keith's solution was sufficient, but I'm still very interested in a solution
> that's independent of dimension and axis.
>   Thanks (and thanks, Keith),
>      Ken

an alternative to Vincent's general solution, if you have unique max
or want all argmax is using a mask

>>> a
array([[ 100.,    0.,    0.],
       [   0.,  100.,  100.],
       [   0.,    0.,    0.]])
>>> b
array([[0, 1, 2],
       [3, 4, 5],
       [6, 7, 8]])

>>> ax=0; b[a==np.expand_dims(a.max(ax),ax)]
array([0, 4, 5])
>>> ax=1; b[a==np.expand_dims(a.max(ax),ax)]
array([0, 4, 5, 6, 7, 8])

>>> aa=np.eye(3)
>>> ax=1; b[aa==np.expand_dims(aa.max(ax),ax)]
array([0, 4, 8])
>>> ax=0; b[aa==np.expand_dims(aa.max(ax),ax)]
array([0, 4, 8])

Josef
>
>
> _______________________________________________
> NumPy-Discussion mailing list
> NumPy-Discussion at scipy.org
> http://mail.scipy.org/mailman/listinfo/numpy-discussion
>
>