Sort each column of an numpy.ndarray using the output of numpy.argsort
As of may 2018 it can done using np.take_along_axis
np.take_along_axis(ref_arr, sm, axis=0)
Out[25]:
array([[10, 16, 15, 10],
[13, 23, 24, 12],
[28, 26, 28, 28]])
Basically two steps are needed :
1] Get the argsort indices along each col with axis=0
-
sidx = ref_arr.argsort(axis=0)
2] Use advanced-indexing
to use sidx
for selecting rows i.e. to index into the first dimension and use another range array to index into the second dimension, so that it would cover sidx
indices across all the columns -
out = ref_arr[sidx, np.arange(sidx.shape[1])]
Sample run -
In [185]: ref_arr
Out[185]:
array([[12, 22, 12, 13],
[28, 26, 21, 23],
[24, 14, 16, 25]])
In [186]: sidx = ref_arr.argsort(axis=0)
In [187]: sidx
Out[187]:
array([[0, 2, 0, 0],
[2, 0, 2, 1],
[1, 1, 1, 2]])
In [188]: ref_arr[sidx, np.arange(sidx.shape[1])]
Out[188]:
array([[12, 14, 12, 13],
[24, 22, 16, 23],
[28, 26, 21, 25]])