Iterar a través de la dimensión arbitraria en numpy

Question

Tengo una matriz numpy multidimensional, y necesito iterar a través de una dimensión determinada. El problema es que no sabré qué dimensión hasta el tiempo de ejecución. En otras palabras, dada una matriz m, que podría desear

m[:,:,:,i] for i in xrange(n) 

o lo que pueda desear

m[:,:,i,:] for i in xrange(n) 

etc.

Imagino que debe haber una función directa en numpy a escribir este , pero no puedo descifrar qué es/cómo se podría llamar. ¿Alguna idea?

Answer 1

Hay muchas formas de hacerlo. Puede construir el índice correcto con una lista de sectores, o quizás alterar los pasos de m. Sin embargo, la forma más sencilla puede ser la de utilizar np.swapaxes:

import numpy as np 
m=np.arange(24).reshape(2,3,4) 
print(m.shape) 
# (2, 3, 4) 

Let axis sea el eje que desea reproducir indefinidamente. m_swapped es lo mismo que m excepto que el eje axis=1 se intercambia con el último eje (axis=-1).

axis=1 
m_swapped=m.swapaxes(axis,-1) 
print(m_swapped.shape) 
# (2, 4, 3) 

Ahora usted puede simplemente bucle sobre el último eje:

for i in xrange(m_swapped.shape[-1]): 
    assert np.all(m[:,i,:] == m_swapped[...,i]) 

Tenga en cuenta que m_swapped es una vista, no una copia, de m. Alterar m_swapped alterará m.

m_swapped[1,2,0]=100 
print(m) 
assert(m[1,0,2]==100) 

Answer 2

Puede utilizar slice(None) en lugar de la :. Por ejemplo,

from numpy import * 

d = 2 # the dimension to iterate 

x = arange(5*5*5).reshape((5,5,5)) 
s = slice(None) # : 

for i in range(5): 
    slicer = [s]*3 # [:, :, :] 
    slicer[d] = i # [:, :, i] 
    print x[slicer] # x[:, :, i]