Itérer sur une dimension arbitraire en numpy

Question

J'ai un tableau numpy multidimensionnel et j'ai besoin d'itérer sur une dimension donnée. Le problème est, je ne saurai pas dimension jusqu'à l'exécution. En d'autres termes, étant donné un m tableau, je pourrais vouloir

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

ou je pourrais vouloir

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

etc.

J'imagine qu'il doit y avoir une fonction simple à numpy d'écrire ce , mais je ne peux pas comprendre ce que c'est/ce que l'on pourrait appeler. Des pensées?

Answer 1

Il y a plusieurs façons de le faire. Vous pouvez créer le bon index avec une liste de tranches, ou peut-être modifier les pas de m. Cependant, la façon la plus simple est peut-être utiliser np.swapaxes:

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

Soit axis être l'axe que vous souhaitez faire une boucle sur. m_swapped correspond à m sauf que l'axe axis=1 est remplacé par le dernier axe (axis=-1).

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

Maintenant, vous pouvez juste boucle sur le dernier axe:

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

Notez que m_swapped est une vue, pas une copie, de m. La modification m_swapped modifie m.

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

Answer 2

Vous pouvez utiliser slice(None) à la place du :. Par exemple,

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]