Riconcilia np.fromiter e matrici multidimensionali in Python

Question

Mi piace usare np.fromiter da numpy perché è un modo risorsa-pigro per creare oggetti np.array. Tuttavia, sembra che non supporti array multidimensionali, che sono anche molto utili.

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 

Come posso ottenere a di essere un array multidimensionale, mantenendo una costruzione come pigro sulla base di generatori?

Answer 1

Di per sé, np.fromiter supporta solo costruire array 1D, e come tale, esso prevede un iterabile che produrrà valori individuali piuttosto che tuple/elenchi/sequenze ecc Un modo per ovviare a questa limitazione sarebbe usare itertools.chain.from_iterable per pigramente 'decomprimere' l'output dell'espressione del generatore in una singola sequenza di valori 1D:

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)