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)
The 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.
If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!
Donate Us With