argsort does not work for multidimensional arrays

[healther]

consider the following code:
a = np.random.random([5,5])
ind = np.argsort(a, axis=1)
a_sorted = a[ind]
np.sort(a, axis=1)

now a_sorted and a should be both sorted along the 1-axis. However not even the shapes are identical anymore. a is still of (5,5) while a_sorted is (5,5,5). If this is the intended behaviour, can someone tell me why?

[rkern]

argsort() is working fine, but indexing doesn't work the way that you are expecting it to.

http://docs.scipy.org/doc/numpy/reference/arrays.indexing.html#advanced-indexing

Try this:

a = np.random.random([5,5])
i = np.arange(len(a))[:, np.newaxis]
j = np.argsort(a, axis=1)
a_sorted = a[i, j]
np.sort(a, axis=1)

[lzkelley]

I was confused in the same way. The documentation specifically states,

Returns
index_array : ndarray, int
Array of indices that sort a along the specified axis.
In other words, a[index_array] yields a sorted a.

Which only applies to a 1D array. Should the documentation be extended to specify that this is not the case for ND arrays?

[shoyer]

@lzkelley Yes, a patch to improve the documentation would be very welcome.

[lzkelley]

@shoyer cool, I'll make a PR for it and update

[endolith]

It would be good if there were an example of how to actually use the output to sort the N-D array.

It would also be good if there were a more readable solution than a[np.arange(np.shape(a)[0])[:,np.newaxis], np.argsort(a)]

[charris]

Maybe @seberg could make it one of his indexing function additions.

[seberg]

I don't think it fits indexing very obviously, a pick function or so is probably an easier match.

[shoyer]

+1 for a new function.

[lzkelley]

Isn't the entire purpose of argsort (as apposed to sort) to be used for indexing?

[endolith]

argsort supplies the indexes to sort along an axis, but actually sorting along that axis using the indexes seems needlessly verbose. Unless there's a simpler way I'm not aware of.

[seberg]

https://github.com/numpy/numpy/blob/5becd12914d0402967205579d6f59a98151e0d98/doc/neps/indexing.rst#related-questions

[eric-wieser]

a pick function or so is probably an easier match.

This is #8708