Reconcile np.fromiter and multidimensional arrays in Python

Question

I love using np.fromiter from numpy because it is a resource-lazy way to build np.array objects. However, it seems like it doesn\'t su

Answer 1

By itself, np.fromiter only supports constructing 1D arrays, and as such, it expects an iterable that will yield individual values rather than tuples/lists/sequences etc. One way to work around this limitation would be to use itertools.chain.from_iterable to lazily 'unpack' the output of your generator expression into a single 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)