2009-10-19 22 views
34

Esiste una funzione per ottenere un iteratore su una dimensione arbitraria di un array numpy?Iterating over arbitrary dimension of numpy.array

iterare sulla prima dimensione è facile ...

In [63]: c = numpy.arange(24).reshape(2,3,4) 

In [64]: for r in c : 
    ....:  print r 
    ....: 
[[ 0 1 2 3] 
[ 4 5 6 7] 
[ 8 9 10 11]] 
[[12 13 14 15] 
[16 17 18 19] 
[20 21 22 23]] 

Ma l'iterazione di altre dimensioni è più difficile. Per esempio, l'ultima dimensione:

In [73]: for r in c.swapaxes(2,0).swapaxes(1,2) : 
    ....:  print r 
    ....: 
[[ 0 4 8] 
[12 16 20]] 
[[ 1 5 9] 
[13 17 21]] 
[[ 2 6 10] 
[14 18 22]] 
[[ 3 7 11] 
[15 19 23]] 

Sto facendo un generatore a fare da solo, ma sono sorpreso che non ci sia una funzione denominata qualcosa come numpy.ndarray.iterdim (asse = 0) a fallo automaticamente.

risposta

36

Quello che proponiamo è abbastanza veloce, ma la leggibilità può essere migliorata con le forme più chiare:

for i in range(c.shape[-1]): 
    print c[:,:,i] 

o, meglio (più veloce, più generale e più esplicito):

for i in range(c.shape[-1]): 
    print c[...,i] 

Tuttavia , il primo approccio sopra sembra essere circa due volte più lento dell'approccio swapaxes():

python -m timeit -s 'import numpy; c = numpy.arange(24).reshape(2,3,4)' \ 
    'for r in c.swapaxes(2,0).swapaxes(1,2): u = r' 
100000 loops, best of 3: 3.69 usec per loop 

python -m timeit -s 'import numpy; c = numpy.arange(24).reshape(2,3,4)' \ 
    'for i in range(c.shape[-1]): u = c[:,:,i]' 
100000 loops, best of 3: 6.08 usec per loop 

python -m timeit -s 'import numpy; c = numpy.arange(24).reshape(2,3,4)' \ 
    'for r in numpy.rollaxis(c, 2): u = r' 
100000 loops, best of 3: 6.46 usec per loop 

Immagino che questo sia dovuto al fatto che swapaxes() non copia alcun dato e perché la gestione di c[:,:,i] potrebbe essere eseguita tramite codice generale (che gestisce il caso in cui : viene sostituito da una porzione più complessa).

Nota, tuttavia, che la seconda soluzione più esplicito c[...,i] è sia abbastanza leggibili e abbastanza veloce:

python -m timeit -s 'import numpy; c = numpy.arange(24).reshape(2,3,4)' \ 
    'for i in range(c.shape[-1]): u = c[...,i]' 
100000 loops, best of 3: 4.74 usec per loop 
3

credo che non v'è alcuna funzione. Quando ho scritto la mia funzione, ho finito per prendere l'iterazione suggerita da EOL. Per i lettori futuri, eccolo:

def iterdim(a, axis=0) : 
    a = numpy.asarray(a); 
    leading_indices = (slice(None),)*axis 
    for i in xrange(a.shape[axis]) : 
    yield a[leading_indices+(i,)] 
+0

La sintassi NumPy standard 'a [..., i]' sarebbe più leggera e rimuoverebbe la necessità di 'leading_indices'. – EOL

+1

@EOL ma funzionerebbe solo per l'ultimo asse, con leading_indices suo più generale ... – lukas

+0

Buon punto @lukas: la domanda iniziale in effetti menziona l'iterazione "su una dimensione arbitraria", mentre pensavo di integrare l'ultima dimensione. – EOL

14

userei il seguente:

c = numpy.arange(2 * 3 * 4) 
c.shape = (2, 3, 4) 

for r in numpy.rollaxis(c, 2): 
    print(r) 

La funzione rollaxis crea una nuova vista sulla matrice. In questo caso sposta l'asse 2 in avanti, equivalente all'operazione c.transpose(2, 0, 1).

+1

+1: abbastanza diretto, ma sfortunatamente un po 'più lento del semplice approccio c [:,:, i] '(non so perché). – EOL

4

Quindi, è possibile scorrere facilmente la prima dimensione, come hai mostrato. Un altro modo per fare ciò per una dimensione arbitraria è usare numpy.rollaxis() per portare la dimensione data al primo (il comportamento predefinito), e quindi usare l'array restituito (che è una vista, quindi è veloce) come un iteratore .

In [1]: array = numpy.arange(24).reshape(2,3,4) 

In [2]: for array_slice in np.rollaxis(array, 1): 
    ....:  print array_slice.shape 
    ....: 
(2, 4) 
(2, 4) 
(2, 4) 

EDIT: io commento che ho presentato un PR per NumPy per affrontare questo qui: https://github.com/numpy/numpy/pull/3262. Il consenso era che questo non era abbastanza da aggiungere al numpy codebase. Penso che l'uso di np.rollaxis sia il modo migliore per farlo, e se vuoi un interatore, avvolgilo in iter().