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.

Paul Panzer · Accepted Answer · 2019-03-09T04:15:53.720

11

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)

edited Mar 09 '19 at 04:15

answered Jan 03 '18 at 06:50

Paul Panzer

51,835
3
54
99

Nice! I didn't think slices would work in this context! – Hameer Abbasi Jan 03 '18 at 06:54
@HameerAbbasi They don't. `ogrid` converts them to `arange`s. – Paul Panzer Jan 03 '18 at 07:12
1

In any case, +1 for making it general. – Hameer Abbasi Jan 03 '18 at 07:39
1

After fiddling a little bit, I successfully implemented your solution into my code. +1 for fully answering my question and giving my a very helpful solution. – 1313e Jan 03 '18 at 08:02
3

You'll want to add `idx = tuple(idx)` now to avoid a `FutureWarning` (and later an error / incorrect result). Something like this is really needed in the `argsort` docs; there should be an easy/standard way to get from the sorted indices to the sorted array. – Nathan Mar 08 '19 at 22:38
1

@Nathan Done, I'm also partially addressing your last point. Maybe `take/put_along_axis` should be linked from the `arg*` function docs. – Paul Panzer Mar 09 '19 at 04:23
1

Thanks for the quick update and adding `take_along_axis`. I notice the docs for `take_along_axis` (and also `put_along_axis`) do reference `argsort` and `argpartition`, so it seems like these are indeed the standard way this should be done now; it would indeed be great if there was a reference in the other direction from the `argsort`/`argprartition` docs. – Nathan Mar 09 '19 at 04:36

hpaulj · Answer 2 · 2018-01-03T07:10:10.993

@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).

This kind of question pops up so frequently, it makes me wonder whether `arg*` functions should perhaps have a flag to make them optionally return all indices needed for getting the non-arg result. — Paul Panzer, Jan 03 '18 at 07:35

score 2 · Answer 3 · answered Jan 03 '18 at 06:45

Here's a vectorized implementation. It should be N-dimensional and quite a bit faster than what you're doing.

import numpy as np


def sort1(array, args):
    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]]

    return array_sort


def sort2(array, args):
    shape = array.shape
    idx = np.ix_(*tuple(np.arange(l) for l in shape[:-1]))
    idx = tuple(ar[..., None] for ar in idx)
    array_sorted = array[idx + (args,)]

    return array_sorted


if __name__ == '__main__':
    array = np.random.rand(5, 6, 7)
    idx = np.argsort(array)

    result1 = sort1(array, idx)
    result2 = sort2(array, idx)

    print(np.array_equal(result1, result2))

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

3 Answers3

Linked