How to use numpy.argsort() as indices in more than 2 dimensions?

Question

I know something similar to this question has been asked many times over already, but all answers given to similar questions only seem to work for arrays with 2 dimensions.

My understanding of np.argsort() is that np.sort(array) == array[np.argsort(array)] should be True. I have found out that this is indeed correct if np.ndim(array) == 2, but it gives different results if np.ndim(array) > 2.

Example:

>>> array = np.array([[[ 0.81774634,  0.62078744],
                       [ 0.43912609,  0.29718462]],
                      [[ 0.1266578 ,  0.82282054],
                       [ 0.98180375,  0.79134389]]])
>>> np.sort(array)
array([[[ 0.62078744,  0.81774634],
        [ 0.29718462,  0.43912609]],

       [[ 0.1266578 ,  0.82282054],
        [ 0.79134389,  0.98180375]]])
>>> array.argsort()
array([[[1, 0],
        [1, 0]],

       [[0, 1],
        [1, 0]]])
>>> array[array.argsort()]
array([[[[[ 0.1266578 ,  0.82282054],
          [ 0.98180375,  0.79134389]],

         [[ 0.81774634,  0.62078744],
          [ 0.43912609,  0.29718462]]],


        [[[ 0.1266578 ,  0.82282054],
          [ 0.98180375,  0.79134389]],

         [[ 0.81774634,  0.62078744],
          [ 0.43912609,  0.29718462]]]],



       [[[[ 0.81774634,  0.62078744],
          [ 0.43912609,  0.29718462]],

         [[ 0.1266578 ,  0.82282054],
          [ 0.98180375,  0.79134389]]],


        [[[ 0.1266578 ,  0.82282054],
          [ 0.98180375,  0.79134389]],

         [[ 0.81774634,  0.62078744],
          [ 0.43912609,  0.29718462]]]]])

So, can anybody explain to me how exactly np.argsort() can be used as the indices to obtain the sorted array? The only way I can come up with is:

args = np.argsort(array)
array_sort = np.zeros_like(array)
for i in range(array.shape[0]):
    for j in range(array.shape[1]):
        array_sort[i, j] = array[i, j, args[i, j]]

which is extremely tedious and cannot be generalized for any given number of dimensions.

Answer 1

Here is a general method:

import numpy as np

array = np.array([[[ 0.81774634,  0.62078744],
                   [ 0.43912609,  0.29718462]],
                  [[ 0.1266578 ,  0.82282054],
                   [ 0.98180375,  0.79134389]]])

a = 1 # or 0 or 2

order = array.argsort(axis=a)

idx = np.ogrid[tuple(map(slice, array.shape))]
# if you don't need full ND generality: in 3D this can be written
# much more readable as
# m, n, k = array.shape
# idx = np.ogrid[:m, :n, :k]

idx[a] = order

print(np.all(array[idx] == np.sort(array, axis=a)))

Output:

True

Explanation: We must specify for each element of the output array the complete index of the corresponding element of the input array. Thus each index into the input array has the same shape as the output array or must be broadcastable to that shape.

The indices for the axes along which we do not sort/argsort stay in place. We therefore need to pass a broadcastable range(array.shape[i]) for each of those. The easiest way is to use ogrid to create such a range for all dimensions (If we used this directly, the array would come back unchanged.) and then replace the index correspondingg to the sort axis with the output of argsort.

UPDATE March 2019:

Numpy is becoming more strict in enforcing multi-axis indices being passed as tuples. Currently, array[idx] will trigger a deprecation warning. To be future proof use array[tuple(idx)] instead. (Thanks @Nathan)

Or use numpy's new (version 1.15.0) convenience function take_along_axis:

np.take_along_axis(array, order, a)

Answer 2

@Hameer's answer works, though it might use some simplification and explanation.

sort and argsort are working on the last axis. argsort returns a 3d array, same shape as the original. The values are the indices on that last axis.

In [17]: np.argsort(arr, axis=2)
Out[17]: 
array([[[1, 0],
        [1, 0]],

       [[0, 1],
        [1, 0]]], dtype=int32)
In [18]: _.shape
Out[18]: (2, 2, 2)
In [19]: idx=np.argsort(arr, axis=2)

To use this we need to construct indices for the other dimensions that broadcast to the same (2,2,2) shape. ix_ is a handy tool for this.

Just using idx as one of the ix_ inputs doesn't work:

In [20]: np.ix_(range(2),range(2),idx)
....
ValueError: Cross index must be 1 dimensional

Instead I use the last range, and then ignore it. @Hameer instead constructs the 2d ix_, and then expands them.

In [21]: I,J,K=np.ix_(range(2),range(2),range(2))
In [22]: arr[I,J,idx]
Out[22]: 
array([[[ 0.62078744,  0.81774634],
        [ 0.29718462,  0.43912609]],

       [[ 0.1266578 ,  0.82282054],
        [ 0.79134389,  0.98180375]]])

So the indices for the other dimensions work with the (2,2,2) idx array:

In [24]: I.shape
Out[24]: (2, 1, 1)
In [25]: J.shape
Out[25]: (1, 2, 1)

That's the basics for constructing the other indices when you are given multidimensional index for one dimension.

@Paul constructs the same indices with ogrid:

In [26]: np.ogrid[slice(2),slice(2),slice(2)]  # np.ogrid[:2,:2,:2]
Out[26]: 
[array([[[0]],

        [[1]]]), array([[[0],
         [1]]]), array([[[0, 1]]])]
In [27]: _[0].shape
Out[27]: (2, 1, 1)

ogrid as a class works with slices, while ix_ requires a list/array/range.

argsort for a multidimensional ndarray (from 2015) works with a 2d array, but the same logic applies (find a range index(s) that broadcasts with the argsort).