Iteration through all 1 dimensional subarrays of a multi-dimensional array

What is the fastest way to iterate through all one dimensional sub-arrays of an n dimensional array in python.

For example consider the 3-D array:

import numpy as np 
a = np.arange(24)
a = a.reshape(2,3,4)

T开发者_Go百科he desired sequence of yields from the iterator is :

a[:,0,0]
a[:,0,1]
..
a[:,2,3]
a[0,:,0]
..
a[1,:,3]
a[0,0,:]
..
a[1,2,:]

---

Here is a compact implementation of such an iterator:

def iter1d(a):
    return itertools.chain.from_iterable(
        numpy.rollaxis(a, axis, a.ndim).reshape(-1, dim)
        for axis, dim in enumerate(a.shape))

This will yield the subarrays in the order you gave in your post:

for x in iter1d(a):
    print x

prints

[ 0 12]
[ 1 13]
[ 2 14]
[ 3 15]
[ 4 16]
[ 5 17]
[ 6 18]
[ 7 19]
[ 8 20]
[ 9 21]
[10 22]
[11 23]
[0 4 8]
[1 5 9]
[ 2  6 10]
[ 3  7 11]
[12 16 20]
[13 17 21]
[14 18 22]
[15 19 23]
[0 1 2 3]
[4 5 6 7]
[ 8  9 10 11]
[12 13 14 15]
[16 17 18 19]
[20 21 22 23]

The trick here is to iterate over all axes, and for each axis reshape the array to a two-dimensional array the rows of which are the desired one-dimensional subarrays.

---

There may be a more efficient way, but this should work...

import itertools
import numpy as np

a = np.arange(24)
a = a.reshape(2,3,4)

colon = slice(None)
dimensions = [range(dim) + [colon] for dim in a.shape]

for dim in itertools.product(*dimensions):
    if dim.count(colon) == 1:
        print a[dim]

This yields (I'm leaving out a trivial bit of code to print the left hand side of this...):

a[0,0,:] -->  [0 1 2 3]
a[0,1,:] -->  [4 5 6 7]
a[0,2,:] -->  [ 8  9 10 11]
a[0,:,0] -->  [0 4 8]
a[0,:,1] -->  [1 5 9]
a[0,:,2] -->  [ 2  6 10]
a[0,:,3] -->  [ 3  7 11]
a[1,0,:] -->  [12 13 14 15]
a[1,1,:] -->  [16 17 18 19]
a[1,2,:] -->  [20 21 22 23]
a[1,:,0] -->  [12 16 20]
a[1,:,1] -->  [13 17 21]
a[1,:,2] -->  [14 18 22]
a[1,:,3] -->  [15 19 23]
a[:,0,0] -->  [ 0 12]
a[:,0,1] -->  [ 1 13]
a[:,0,2] -->  [ 2 14]
a[:,0,3] -->  [ 3 15]
a[:,1,0] -->  [ 4 16]
a[:,1,1] -->  [ 5 17]
a[:,1,2] -->  [ 6 18]
a[:,1,3] -->  [ 7 19]
a[:,2,0] -->  [ 8 20]
a[:,2,1] -->  [ 9 21]
a[:,2,2] -->  [10 22]
a[:,2,3] -->  [11 23]

The key here is that indexing a with (for example) a[0,0,:] is equivalent to indexing a with a[(0,0,slice(None))]. (This is just generic python slicing, nothing numpy-specific. To prove it to yourself, you can write a dummy class with just a __getitem__ and print what's passed in when you index an instance of your dummy class.).

So, what we want is every possible combination of 0 to nx, 0 to ny, 0 to nz, etc and a None for each axis.

However, we want 1D arrays, so we need to filter out anything with more or less than one None (i.e. we don't want a[:,:,:], a[0,:,:], a[0,0,0] etc).

Hopefully that makes some sense, anyway...

Edit: I'm assuming that the exact order doesn't matter... If you need the exact ordering you list in your question, you'll need to modify this...

---

Your friends are the slice() objects, numpy's ndarray.__getitem__() method, and possibly the itertools.chain.from_iterable.