Igual aesta resposta, Tenho um par de matrizes numpy 3D,a eb e quero classificar as entradas deb pelos valores dea. Ao contrário deesta resposta, Quero classificar apenas ao longo de um eixo das matrize

Minha leitura ingênua donumpy.argsort() documentação:

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`.

me levou a acreditar que eu poderia fazer minha classificação com o seguinte código:

import numpy

a = numpy.zeros((3, 3, 3))
a += numpy.array((1, 3, 2)).reshape((3, 1, 1))
print "a"
print a
"""
[[[ 1.  1.  1.]
  [ 1.  1.  1.]
  [ 1.  1.  1.]]

 [[ 3.  3.  3.]
  [ 3.  3.  3.]
  [ 3.  3.  3.]]

 [[ 2.  2.  2.]
  [ 2.  2.  2.]
  [ 2.  2.  2.]]]
"""
b = numpy.arange(3*3*3).reshape((3, 3, 3))
print "b"
print b
"""
[[[ 0  1  2]
  [ 3  4  5]
  [ 6  7  8]]

 [[ 9 10 11]
  [12 13 14]
  [15 16 17]]

 [[18 19 20]
  [21 22 23]
  [24 25 26]]]
"""
print "a, sorted"
print numpy.sort(a, axis=0)
"""
[[[ 1.  1.  1.]
  [ 1.  1.  1.]
  [ 1.  1.  1.]]

 [[ 2.  2.  2.]
  [ 2.  2.  2.]
  [ 2.  2.  2.]]

 [[ 3.  3.  3.]
  [ 3.  3.  3.]
  [ 3.  3.  3.]]]
"""

##This isnt' working how I'd like
sort_indices = numpy.argsort(a, axis=0)
c = b[sort_indices]
"""
Desired output:

[[[ 0  1  2]
  [ 3  4  5]
  [ 6  7  8]]

 [[18 19 20]
  [21 22 23]
  [24 25 26]]

 [[ 9 10 11]
  [12 13 14]
  [15 16 17]]]
"""
print "Desired shape of b[sort_indices]: (3, 3, 3)."
print "Actual shape of b[sort_indices]:"
print c.shape
"""
(3, 3, 3, 3, 3)
"""

Qual é o caminho certo para fazer isso?