Reconcile np.fromiter and multidimensional arrays in Python

Question

I like to use np.fromiter from numpy because it is a resource-lazy way to create np.array objects. However, it does not seem to support multidimensional arrays, which are also quite useful.

 import numpy as np def fun(i): """ A function returning 4 values of the same type. """ return tuple(4*i + j for j in range(4)) # Trying to create a 2-dimensional array from it: a = np.fromiter((fun(i) for i in range(5)), '4i', 5) # fails # This function only seems to work for 1D array, trying then: a = np.fromiter((fun(i) for i in range(5)), [('', 'i'), ('', 'i'), ('', 'i'), ('', 'i')], 5) # painful # .. `a` now looks like a 2D array but it is not: a.transpose() # doesn't work as expected a[0, 1] # too many indices (of course) a[:, 1] # don't even think about it 

How can I get a as a multidimensional array while preserving such a lazy generator-based construct?

Answer 1

By itself, np.fromiter only supports the construction of 1D arrays and, as such, expects iterability to give separate values, not tuples / lists / sequence, etc. One way around this limitation would be to use itertools.chain.from_iterable to lazily "unzip" the output of your generator expression into one 1D sequence of values:

 import numpy as np from itertools import chain def fun(i): return tuple(4*i + j for j in range(4)) a = np.fromiter(chain.from_iterable(fun(i) for i in range(5)), 'i', 5 * 4) a.shape = 5, 4 print(repr(a)) # array([[ 0, 1, 2, 3], # [ 4, 5, 6, 7], # [ 8, 9, 10, 11], # [12, 13, 14, 15], # [16, 17, 18, 19]], dtype=int32)