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

Question

similares a this answer, tengo un par de paneles 3D numpy, a y b, y quiero ordenar las entradas de b por los valores de a. A diferencia de this answer, quiero ordenar solo a lo largo de un eje de las matrices.

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

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 llevó a creer que podía hacer mi clase 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?

Answer 1

Aún debe 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 "aplanado" a través de los otros dos ejes (o n - 1 ejes, donde n = número total de ejes). En otras palabras, este (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 para 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, podría utilizar numpy.ogrid, como sugiere unutbu. Como el objeto generado por ogrid es más pequeño, crearé los tres ejes, solo por consistencia, pero tenga en cuenta el comentario de unutbu sobre cómo 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.]]])