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

Question

我有一个类似下面的迭代器

it = ((x, x**2) for x in range(20))

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

如果我只有一个大小未知的结果，我可以使用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 无论如何都会填充整个列表大小的双端队列。我很想避免将迭代器复制到列表中，然后再将其再次复制到数组中，但我想不出一种方法来增量构建数组而不完全手动完成。另外，fromiter 好像是用 c 写的，所以用 python 写可能和先做一个列表没什么区别。

Answer 1

您可以使用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

因此您可能希望使y 和y2 连续：

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