¿Cómo puedo obtener elementos de la matriz 3D usando índices específicos en numpy?
- Tengo una matriz 3D, en el siguiente ejemplo es una matriz (5, 4, 2):
data_matrix
- Tengo otra matriz de índice de forma (5, 4) donde cada fila de la matriz representa la ubicación del elemento:
indx_array
No sé cómo puedo conseguir el required_output
. Estoy tratando de organizar (1,2) elementos de cada fila según elindx_array
¡No quiero usar bucles for!
data_matrix = np.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, 27], [28, 29], [30, 31]],
[[32, 33], [34, 35], [36, 37], [38, 39]]
])
indx_array = np.array([[3,2,1,0], [0,1,2,3], [1,0,3,2], [0,3,1,2], [1,2,3,0]])
# I want following result:
required_output = [
[[6, 7], [4, 5], [2, 3], [0, 1]]
[[8, 9], [10, 11], [12, 13], [14, 15]]
[[18, 19], [16, 17], [22, 23], [20, 21]]
[[24, 25], [30, 31], [26, 27], [28, 29]]
[[34, 35], [36, 37], [38, 39], [32, 33]]
]
EDITAR: actualizado indx_array
para ilustrar mejor la situación.
Respuestas
- Numpy: Indexación
- Numpy: Indexación de matrices multidimensionales
In [637]: data_matrix.shape
Out[637]: (5, 4, 2)
In [638]: indx_array.shape
Out[638]: (5, 4)
Necesita indexación avanzada en las dos primeras dimensiones. La matriz de la primera dimensión debe transmitirse con la segunda (5,4). Para hacer eso hago un (5,1) arange
:
In [639]: data_matrix[np.arange(5)[:,None], indx_array]
Out[639]:
array([[[ 6, 7],
[ 4, 5],
[ 2, 3],
[ 0, 1]],
[[ 8, 9],
[10, 11],
[12, 13],
[14, 15]],
[[18, 19],
[16, 17],
[22, 23],
[20, 21]],
[[24, 25],
[30, 31],
[26, 27],
[28, 29]],
[[34, 35],
[36, 37],
[38, 39],
[32, 33]]])
Contraste mi índice (5,1) con el aceptado _x
(que es (5,4) entrelazado):
In [640]: np.arange(5)[:,None]
Out[640]:
array([[0],
[1],
[2],
[3],
[4]])
In [641]: _x = np.repeat(np.arange(indx_array.shape[0]),indx_array.shape[1])
In [643]: _x
Out[643]: array([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4])
Con la radiodifusión _x
, no necesita repetirse, (5,4); (5,1) es suficiente.
La radiodifusión hace una repetición virtual. Esto se puede ilustrar con la broadcast_to
función:
In [648]: np.broadcast_to(np.arange(5)[:,None],(5,4))
Out[648]:
array([[0, 0, 0, 0],
[1, 1, 1, 1],
[2, 2, 2, 2],
[3, 3, 3, 3],
[4, 4, 4, 4]])
In [649]: _.strides
Out[649]: (8, 0)
Es ese paso el 0
que se repite sin hacer copias. as_strided
es la stride_tricks
función más útil , especialmente cuando se hacen cosas como mover ventanas. Por lo general, dejamos que la transmisión automática haga el trabajo sin preocuparnos demasiado por el cómo.
Se puede hacer con un poco de manipulación de la matriz de índices.
import numpy as np
_x = np.repeat(np.arange(indx_array.shape[0]),indx_array.shape[1])
_y = indx_array.ravel()
output = data_matrix[_x, _y].reshape(data_matrix.shape)
lo que da como resultado la matriz numpy esperada
array([[[ 6, 7],
[ 4, 5],
[ 2, 3],
[ 0, 1]],
[[ 8, 9],
[10, 11],
[12, 13],
[14, 15]],
[[18, 19],
[16, 17],
[22, 23],
[20, 21]],
[[24, 25],
[30, 31],
[26, 27],
[28, 29]],
[[34, 35],
[36, 37],
[38, 39],
[32, 33]]])