numpy fromiter with generator of list

前端 未结 2 1022
佛祖请我去吃肉
佛祖请我去吃肉 2021-01-12 09:03
import numpy as np
def gen_c():
    c = np.ones(5, dtype=int)
    j = 0
    t = 10
    while j < t:
        c[0] = j
        yield c.tolist()
        j += 1 

# W         


        
2条回答
  •  抹茶落季
    2021-01-12 09:33

    As the docs suggested, np.fromiter() only accepts 1-dimensional iterables. You can use itertools.chain.from_iterable() to flatten the iterable first, and np.reshape() it back later:

    import itertools
    import numpy as np
    
    def fromiter2d(it, dtype):
    
        # clone the iterator to get its length
        it, it2 = itertools.tee(it)
        length = sum(1 for _ in it2)
    
        flattened = itertools.chain.from_iterable(it)
        array_1d = np.fromiter(flattened, dtype)
        array_2d = np.reshape(array_1d, (length, -1))
        return array_2d
    

    Demo:

    >>> iter2d = (range(i, i + 4) for i in range(0, 12, 4))
    
    >>> from_2d_iter(iter2d, int)
    array([[ 0,  1,  2,  3],
           [ 4,  5,  6,  7],
           [ 8,  9, 10, 11]])
    

    Only tested on Python 3.6, but should also work with Python 2.

提交回复
热议问题