Acceso n-th dimensión en python

Quiero un acceso fácil de leer a algunas partes de una matriz numpy multidimensional. Para cualquier matriz, acceder a la primera dimensión es fácil ( b[index] ). Acceder a la sexta dimensión por otro lado es “difícil” (especialmente para leer).

 b[:,:,:,:,:,index] #the next person to read the code will have to count the : 

¿Hay una mejor manera de hacer esto? Especialmente, ¿hay alguna forma de que no se conozca el eje al escribir el progtwig?

Edición: la dimensión indexada no es necesariamente la última dimensión

Si desea una vista y la desea rápidamente, puede crear el índice manualmente:

 arr[(slice(None), )*5 + (your_index, )] # ^---- This is equivalent to 5 colons: `:, :, :, :, :` 

Lo cual es mucho más rápido que np.take y solo ligeramente más lento que la indexación con : s:

 import numpy as np arr = np.random.random((10, 10, 10, 10, 10, 10, 10)) np.testing.assert_array_equal(arr[:,:,:,:,:,4], arr.take(4, axis=5)) np.testing.assert_array_equal(arr[:,:,:,:,:,4], arr[(slice(None), )*5 + (4, )]) %timeit arr.take(4, axis=5) # 18.6 ms ± 249 µs per loop (mean ± std. dev. of 7 runs, 10 loops each) %timeit arr[(slice(None), )*5 + (4, )] # 2.72 µs ± 39.7 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each) %timeit arr[:, :, :, :, :, 4] # 2.29 µs ± 107 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each) 

Pero tal vez no sea tan legible, por lo que si necesita eso a menudo, probablemente debería ponerlo en una función con un nombre significativo:

 def index_axis(arr, index, axis): return arr[(slice(None), )*axis + (index, )] np.testing.assert_array_equal(arr[:,:,:,:,:,4], index_axis(arr, 4, axis=5)) %timeit index_axis(arr, 4, axis=5) # 3.79 µs ± 127 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each) 

Puedes usar np.take . Por ejemplo:

 b.take(index, axis=5) 

Una forma intermedia (en legibilidad y tiempo) entre las respuestas de MSeifert y kazemakase es usar np.rollaxis :

 np.rollaxis(b, axis=5)[index] 

Probando las soluciones:

 import numpy as np arr = np.random.random((10, 10, 10, 10, 10, 10, 10)) np.testing.assert_array_equal(arr[:,:,:,:,:,4], arr.take(4, axis=5)) np.testing.assert_array_equal(arr[:,:,:,:,:,4], arr[(slice(None), )*5 + (4, )]) np.testing.assert_array_equal(arr[:,:,:,:,:,4], np.rollaxis(arr, 5)[4]) %timeit arr.take(4, axis=5) # 100 loops, best of 3: 4.44 ms per loop %timeit arr[(slice(None), )*5 + (4, )] # 1000000 loops, best of 3: 731 ns per loop %timeit arr[:, :, :, :, :, 4] # 1000000 loops, best of 3: 540 ns per loop %timeit np.rollaxis(arr, 5)[4] # 100000 loops, best of 3: 3.41 µs per loop 

En el espíritu de la rollaxis de @ Jürg Merlin Spaak pero mucho más rápido y no en desuso :

 b.swapaxes(0, axis)[index] 

Puedes decir:

 slice = b[..., index]