Numpy同时使用fromiter创建两个数组

时间:2016-02-04 21:47:32

标签: python arrays numpy

我有一个类似于下面的

的迭代器
it = ((x, x**2) for x in range(20))

我想要的是两个数组。其中一个x和另一个x**2但我实际上并不知道元素的数量,我无法从一个条目转换到另一个条目,所以我不能构建第一个,然后从第一个构建第二个。

如果我只有一个未知大小的结果,我可以使用np.fromiter来有效地动态分配,例如。

y = np.fromiter((x[0] for x in it), float)

有两个我希望我能做点像

ita, itb = itertools.tee(it)
y = np.fromiter((x[0] for x in ita), float)
y2 = np.fromiter((x[1] for x in itb), float)

但是因为第一次调用耗尽了迭代器,所以我最好不要做

lst = list(it)
y = np.fromiter((x[0] for x in lst), float, len(lst))
y2 = np.fromiter((x[1] for x in lst), float, len(lst))

因为tee将填充整个列表大小的deque。我希望避免将迭代器复制到列表中,然后再将其复制到数组中,但我想不出一种方法来逐步构建数组而不完全手动完成。另外,fromiter似乎是用c编写的,所以在python中编写它可能最终会比首先列出一个列表没有微不足道的差别。

1 个答案:

答案 0 :(得分:2)

您可以使用np.fromiter构建一个包含所有值的数组,然后对数组进行切片:

In [103]: it = ((x, x**2) for x in range(20))

In [104]: import itertools

In [105]: y = np.fromiter(itertools.chain.from_iterable(it), dtype=float)

In [106]: y
Out[106]: 
array([   0.,    0.,    1.,    1.,    2.,    4.,    3.,    9.,    4.,
         16.,    5.,   25.,    6.,   36.,    7.,   49.,    8.,   64.,
          9.,   81.,   10.,  100.,   11.,  121.,   12.,  144.,   13.,
        169.,   14.,  196.,   15.,  225.,   16.,  256.,   17.,  289.,
         18.,  324.,   19.,  361.])

In [107]: y, y2 = y[::2], y[1::2]

In [108]: y
Out[108]: 
array([  0.,   1.,   2.,   3.,   4.,   5.,   6.,   7.,   8.,   9.,  10.,
        11.,  12.,  13.,  14.,  15.,  16.,  17.,  18.,  19.])

In [109]: y2
Out[109]: 
array([   0.,    1.,    4.,    9.,   16.,   25.,   36.,   49.,   64.,
         81.,  100.,  121.,  144.,  169.,  196.,  225.,  256.,  289.,
        324.,  361.])

上面设法将迭代器中的数据加载到数组中而不使用中间Python列表。但是,数组中的基础数据不是连续的。连续数组上的许多操作都更快:

In [19]: a = np.arange(10**6)

In [20]: y1 = a[::2]

In [21]: z1 = np.ascontiguousarray(y1)

In [24]: %timeit y1.sum()
1000 loops, best of 3: 975 µs per loop

In [25]: %timeit z1.sum()
1000 loops, best of 3: 464 µs per loop

所以你可能希望yy2连续:

y = np.ascontiguousarray(y)
y2 = np.ascontiguousarray(y2)

调用np.ascontiguousarray需要复制y中的非连续数据 和y2到新数组。不幸的是,我没有办法创建y和。{ y2为没有复制的连续数组。

这是一个比较使用中间Python列表和NumPy切片(有和没有ascontiguousarray)的基准:

import numpy as np
import itertools as IT

def using_intermediate_list(g):
    lst = list(g)
    y = np.fromiter((x[0] for x in lst), float, len(lst))
    y2 = np.fromiter((x[1] for x in lst), float, len(lst))
    return y, y2

def using_slices(g):
    y = np.fromiter(IT.chain.from_iterable(g), dtype=float)
    y, y2 = y[::2], y[1::2]
    return y, y2

def using_slices_contiguous(g):
    y = np.fromiter(IT.chain.from_iterable(g), dtype=float)
    y, y2 = y[::2], y[1::2]
    y = np.ascontiguousarray(y)
    y2 = np.ascontiguousarray(y2)
    return y, y2

def using_array(g):
    y = np.array(list(g))
    y, y2 = y[:, 0], y[:, 1]
    return y, y2
In [27]: %timeit using_intermediate_list(((x, x**2) for x in range(10**6)))
1 loops, best of 3: 376 ms per loop

In [28]: %timeit using_slices(((x, x**2) for x in range(10**6)))
1 loops, best of 3: 220 ms per loop

In [29]: %timeit using_slices_contiguous(((x, x**2) for x in range(10**6)))
1 loops, best of 3: 237 ms per loop

In [34]: %timeit using_array(((x, x**2) for x in range(10**6)))
1 loops, best of 3: 707 ms per loop