[Numpy-discussion] Strange behavior for argsort() and take()

[Anne Archibald]

2008/6/18 Stéfan van der Walt <stefan at sun.ac.za>:
> 2008/6/18 Anne Archibald <peridot.faceted at gmail.com>:
>> In [7]: x.take(x.argsort())
>> Out[7]: array([ 0. ,  0.1,  0.2,  0.3])
>>
>> If you would like to think of it more mathematically, when you feed
>> np.argsort() an array that represents a permutation of the numbers
>> 0,1,...,n-1, you get back the inverse permutation. When you pass a
>> permutation as the argument to x.take(), you apply the permutation to
>> x. (You can also invert a permutation by feeding it as indices to
>> arange(n)).
>>
>> I have been tempted to write some support functions for manipulating
>> permutations, but I'm not sure how generally useful they would be.
>
> Do we have an easy way of grabbing the elements out of an array that
> correspond to
>
> argmax(x, axis) ?
>
> `take` won't do the job; there `axis` has a different meaning.

Well, here's a quick hack:

def take_axis(X, ix, axis):
    XX = np.rollaxis(X,axis)
    s = XX.shape
    return XX[(ix,)+tuple(np.indices(s[1:]))]

And a version that works for argsort()ing:

def take_axis_argsort(X, ix, axis):
    XX = np.rollaxis(X,axis)
    ixix = np.rollaxis(ix,axis)
    s = XX.shape
    return np.rollaxis(XX[(ixix,)+tuple(np.indices(s)[1:])],0,axis+1)

I'm always a little bit leery of using indices(), since it can produce
very large arrays:

In [60]: np.indices((2,3,4)).shape
Out[60]: (3, 2, 3, 4)

Anne

P.S. rollaxis() saves so very much pain; I do wonder if "unrollaxis"
would be useful, even though it's just
def unrollaxis(X, axis):
    return rollaxis(X, 0, axis+1)
-A