python中itertools.tee的替代品

时间:2015-10-14 22:47:41

标签: python itertools

我正在对几个文件中的大量输入数据进行处理。尝试将处理算法与I / O分开,我使用生成器设置了所有内容。这非常有效,除非我想对通过生成器的数据进行一些中间操作。这是一个点击重点的例子

import numpy as np
from itertools import izip, tee

# Have two input matrices.  In reality they're very large, so data is provided
# one row at a time via generators.
M, N = 100, 3
def gen_data_rows(m,n):
    for i in range(m):
        yield np.random.normal(size=n)

rows1 = gen_data_rows(M,N)
rows2 = gen_data_rows(M,N)

# Signal processing operates on chunks of the input, e.g. blocks of rows and
# yields results at a reduced rate.  Here's a simple example.
def foo_rows(rows):
    i = 0
    for row in rows:
        if i % 5 == 0:
            yield row
        i += 1

# But what if we want to do some transformations between the raw input data
# and the processing?
def fun1(x, y):
    return x + y

def fun2(x, y):
    return (x + y)**2

def foo_transformed_rows(rows1, rows2):
    # Define a generator that consumes both inputs at the same time and
    # produces two streams of output I'd like to send to foo_rows().
    def gen_transformed_rows(rows1, rows2):
        for x, y in izip(rows1, rows2):
            yield fun1(x,y), fun2(x,y)

    # Do I really need to tee the above and define separate generators to pick
    # off each result?
    def pick_generator_idx(gen, i):
        for vals in gen:
            yield vals[i]

    gen_xformed_rows, dupe = tee(gen_transformed_rows(rows1, rows2))
    gen_foo_fun1 = foo_rows(pick_generator_idx(gen_xformed_rows, 0))
    gen_foo_fun2 = foo_rows(pick_generator_idx(dupe, 1))
    for foo1, foo2 in izip(gen_foo_fun1, gen_foo_fun2):
        yield foo1, foo2


for foo1, foo2 in foo_transformed_rows(rows1, rows2):
    print foo1, foo2

我认为这里的主要复杂因素是我有两个输入,我想要组合成两个中间生成器(I / O是瓶颈,所以我真的不想两次运行数据)。有没有更好的方法来实现foo_transformed_rows()功能?必须tee()所需数据并定义生成器只是为了从元组中选择项目似乎有点过分。

编辑:我稍微修改了示例以回应评论,但不幸的是,为了保持完整,它仍然很长。基本问题是处理多输入多输出(MIMO)数据流。我想我喜欢像生成多个生成器的yield语句,例如

def two_streams(gen_a, gen_b):
    "Consumes two generators, produces two results."
    for a, b in itertools.izip(gen_a, gen_b):
        c, d = foo(a, b)
        yield c, d

# This doesn't work.  You get one generator of tuples instead of
# two generators of singletons.
gen_c, gen_d = two_streams(gen_a, gen_b)

我想也许会有一些itertools魔法来做相同的事情。

1 个答案:

答案 0 :(得分:1)

我同意@ ShadowRanger的评论,我不明白你为什么要避开tee。它适用于此目的。

但是,对于我来说,开发原始发电机似乎更简单,更直观:

def transform_rows(fun, rows1, rows2):
    for x, y in izip(rows1, rows2):
        yield fun(x,y)

rows1a, rows1b = tee(rows1)
rows2a, rows2b = tee(rows2)
gen_foo_fun1 = foo_rows(transform_rows(fun1, rows1a, rows2a)
gen_foo_fun2 = foo_rows(transform_rows(fun2, rows1b, rows2b)