Ordenar una matriz numpy por otra matriz, a lo largo de un eje particular

Al igual que en esta respuesta , tengo un par de matrices numpy 3D, b , y quiero ordenar las entradas de b por los valores de a . A diferencia de esta respuesta , quiero ordenar solo a lo largo de un eje de las matrices.

Mi lectura ingenua de la documentación numpy.argsort() :

 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 hizo creer que podía hacer mi orden con el siguiente 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) """ 

¿Cuál es la forma correcta de hacer esto?

Aún tiene que proporcionar índices para las otras dos dimensiones para que esto funcione correctamente.

 >>> a = numpy.zeros((3, 3, 3)) >>> a += numpy.array((1, 3, 2)).reshape((3, 1, 1)) >>> b = numpy.arange(3*3*3).reshape((3, 3, 3)) >>> sort_indices = numpy.argsort(a, axis=0) >>> static_indices = numpy.indices((3, 3, 3)) >>> b[sort_indices, static_indices[1], static_indices[2]] array([[[ 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]]]) 

numpy.indices calcula los índices de cada eje de la matriz cuando se “aplana” a través de los otros dos ejes (o n – 1 ejes donde n = número total de ejes). En otras palabras, esto (disculpas por el largo post):

 >>> static_indices array([[[[0, 0, 0], [0, 0, 0], [0, 0, 0]], [[1, 1, 1], [1, 1, 1], [1, 1, 1]], [[2, 2, 2], [2, 2, 2], [2, 2, 2]]], [[[0, 0, 0], [1, 1, 1], [2, 2, 2]], [[0, 0, 0], [1, 1, 1], [2, 2, 2]], [[0, 0, 0], [1, 1, 1], [2, 2, 2]]], [[[0, 1, 2], [0, 1, 2], [0, 1, 2]], [[0, 1, 2], [0, 1, 2], [0, 1, 2]], [[0, 1, 2], [0, 1, 2], [0, 1, 2]]]]) 

Estos son los índices de identidad de cada eje; cuando se usan para indexar b, recrean b.

 >>> b[static_indices[0], static_indices[1], static_indices[2]] array([[[ 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]]]) 

Como alternativa a numpy.indices , puede usar numpy.ogrid , como sugiere unutbu. Dado que el objeto generado por ogrid es más pequeño, crearé los tres ejes, solo por coherencia, pero tenga en cuenta el comentario de unutbu para una forma de hacerlo generando solo dos.

 >>> static_indices = numpy.ogrid[0:a.shape[0], 0:a.shape[1], 0:a.shape[2]] >>> a[sort_indices, static_indices[1], static_indices[2]] array([[[ 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.]]])