Numpy: ordenar una matriz multidimensional por una matriz multidimensional

Question

Perdóname si esto es redundante o súper básico. Voy a Python/Numpy de R y me cuesta mucho cambiar las cosas en mi cabeza.

Tengo una matriz n dimensional que quiero ordenar usando otra matriz n dimensional de valores de índice. Sé que podría envolver esto en un ciclo, pero parece que debería haber una forma Numpyonic realmente concisa de llevar esto a sumisión. Aquí está mi código de ejemplo para establecer el problema donde n = 2:

a1 = random.standard_normal(size=[2,5]) 
index = array([[0,1,2,4,3] , [0,1,2,3,4] ]) 

por lo que ahora tiene una matriz de 2 x 5 de números aleatorios y un índice de 2 x 5. He leído la ayuda para take() unas 10 veces, pero mi cerebro no está arruinándolo, obviamente.

pensé que esto podría conseguirme allí:

take(a1, index) 

array([[ 0.29589188, -0.71279375, -0.18154864, -1.12184984, 0.25698875], 
     [ 0.29589188, -0.71279375, -0.18154864, 0.25698875, -1.12184984]]) 

pero eso es claramente reordenamiento sólo el primer elemento (supongo que a causa de aplanamiento).

Cualquier consejo sobre cómo llego desde donde estoy a una solución que ordena el elemento 0 de a1 por el elemento 0 del índice ... elemento n?

Answer 1

No puedo pensar en cómo trabajar esto en n dimensiones todavía, pero aquí es la versión 2D:

>>> a = np.random.standard_normal(size=(2,5)) 
>>> a 
array([[ 0.72322499, -0.05376714, -0.28316358, 1.43025844, -0.90814293], 
     [ 0.7459107 , 0.43020728, 0.05411805, -0.32813465, 2.38829386]]) 
>>> i = np.array([[0,1,2,4,3],[0,1,2,3,4]]) 
>>> a[np.arange(a.shape[0])[:,np.newaxis],i] 
array([[ 0.72322499, -0.05376714, -0.28316358, -0.90814293, 1.43025844], 
     [ 0.7459107 , 0.43020728, 0.05411805, -0.32813465, 2.38829386]]) 

Aquí está la versión de N dimensiones:

>>> a[list(np.ogrid[[slice(x) for x in a.shape]][:-1])+[i]] 

Así es como funciona:

Bien, empecemos con una matriz tridimensional para ilustración.

>>> import numpy as np 
>>> a = np.arange(24).reshape((2,3,4)) 
>>> a 
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]]]) 

Se puede acceder a elementos de esta matriz especificando el índice a lo largo de cada eje de la siguiente manera:

>>> a[0,1,2] 
6 

Esto es equivalente a a[0][1][2] que es la forma que accedería a un mismo elemento si se tratara de una lista en lugar de una matriz.

Numpy le permite obtener incluso más elegante al rebanar matrices:

>>> a[[0,1],[1,1],[2,2]] 
array([ 6, 18]) 
>>> a[[0,1],[1,2],[2,2]] 
array([ 6, 22]) 

Estos ejemplos serían equivalentes a [a[0][1][2],a[1][1][2]] y [a[0][1][2],a[1][2][2]] si se trata de listas.

Incluso puede omitir los índices repetidos y Numpy descubrirá lo que quiere. Por ejemplo, los ejemplos anteriores podrían ser equivalente escrito:

>>> a[[0,1],1,2] 
array([ 6, 18]) 
>>> a[[0,1],[1,2],2] 
array([ 6, 22]) 

La forma de la matriz (o lista) que se mire con en cada dimensión sólo afecta al forma de la matriz devuelta. En otras palabras, a numpy no le importa que intente indexar su matriz con una matriz de forma (2,3,4) cuando está extrayendo valores, excepto que le devolverá una matriz de forma (2,3,4).Por ejemplo:

>>> a[[[0,0],[0,0]],[[0,0],[0,0]],[[0,0],[0,0]]] 
array([[0, 0], 
     [0, 0]]) 

En este caso, estamos agarrando el mismo elemento, a[0,0,0] una y otra vez, pero numpy está volviendo una matriz con la misma forma que pasamos en

Ok, en. tu problema. Lo que desea es indexar la matriz a lo largo del último eje con los números en su matriz index. Entonces, para el ejemplo en su pregunta, quisiera [[a[0,0],a[0,1],a[0,2],a[0,4],a[0,3]],a[1,0],a[1,1],...

El hecho de que su matriz de índices sea multidimensional, como dije antes, no dice nada acerca de dónde desea extraer estos índices; solo especifica la forma de la matriz de salida. Por lo tanto, en su ejemplo, debe decir numpy que los primeros 5 valores se extraerán de a[0] y los últimos 5 de a[1]. ¡Fácil!

>>> a[[[0]*5,[1]*5],index] 

se complica en N dimensiones, pero vamos a hacerlo para el 3 matriz bidimensional a definí muy por encima. Supongamos que tenemos la matriz de índice siguiente:

>>> i = np.array(range(4)[::-1]*6).reshape(a.shape) 
>>> i 
array([[[3, 2, 1, 0], 
     [3, 2, 1, 0], 
     [3, 2, 1, 0]], 

     [[3, 2, 1, 0], 
     [3, 2, 1, 0], 
     [3, 2, 1, 0]]]) 

Por lo tanto, estos valores son todos para índices a lo largo del último eje. Necesitamos saber qué índices a lo largo del primer y segundo eje se tomarán de estos números; es decir, tenemos que decirle numpy que los índices de los primeros ejes son:

i1 = [[[0, 0, 0, 0], 
     [0, 0, 0, 0], 
     [0, 0, 0, 0]], 

     [[1, 1, 1, 1], 
     [1, 1, 1, 1], 
     [1, 1, 1, 1]]] 

y los índices para el segundo eje son:

i2 = [[[0, 0, 0, 0], 
     [1, 1, 1, 1], 
     [2, 2, 2, 2]], 

     [[0, 0, 0, 0], 
     [1, 1, 1, 1], 
     [2, 2, 2, 2]]] 

Entonces sólo podemos hacer:

>>> a[i1,i2,i] 
array([[[ 3, 2, 1, 0], 
     [ 7, 6, 5, 4], 
     [11, 10, 9, 8]], 

     [[15, 14, 13, 12], 
     [19, 18, 17, 16], 
     [23, 22, 21, 20]]]) 

La útil función numpy que genera i1 y i2 se llama np.mgrid. Uso np.ogrid en mi respuesta, que es equivalente en este caso debido a la magia numpy de la que hablé antes.

Espero que ayude!

Answer 2

Después de jugar con esto un poco más hoy me di cuenta de que si utiliza una función de asignador junto con la toma que podía resolver la versión 2 dimensiones muy simple como esto:

a1 = random.standard_normal(size=[2,5]) 
index = array([[0,1,2,4,3] , [0,1,2,3,4] ]) 
map(take, a1, index) 

Necesitaba map() la take() a cada elemento en a1

Por supuesto, la respuesta aceptada resuelve la versión n-dimensional. Sin embargo, en retrospectiva, determiné que realmente no necesito la solución n-dimensional, solo la versión 2-D.