替换NumPy中的迭代

Question

我正在尝试重写以下例程(a,b,c,d和e都是数组):

def generate_a(b, c, e):
    a = np.zeros_like(b)
    d = np.zeros_like(b)
    for i in range(a.size):
        a[i] = (b[i] / c[i]) + d[i-1]
        d[i] = a[i] * e[i]
    return a

不使用'for'循环,因为例程需要运行数百万次.这是一个小技巧,因为d中任何给定单元格的值取决于为a中的单元格计算的结果,而该结果又取决于为d计算的最后一个值.关于如何在没有迭代的情况下完成此任务的任何想法？

Answer 1

受此启发,如果可以跟踪递归solution,这是一个显示递归vectorizable的帖子.作为链接的解决方案规定,这同样是不应该要快,截至目前,但并行处理器充分利用,这种方法可以给予了一枪.因此,请将此帖作为如何跟踪和矢量化递归的指导.这是实施 -

def generate_a_vectorized(b, c, e):

    K = (b/c)*e
    N = e.size

    mask = np.linspace(N,1,N,dtype=int)[:,None] <= np.arange(N)

    Pn = np.tile(e[None],(N,1))
    Pn[mask] = 1
    En = Pn[:,::-1].cumprod(1)[:,::-1]
    En[mask] = 0
    An = np.append(0,K[:-1])

    d_out = En.dot(An)[::-1] + K
    return (b/c) + np.append(0,d_out[:-1])

样品运行和验证输出 -

In [279]: M = 50
     ...: b = np.random.rand(M)
     ...: c = np.random.rand(M)
     ...: e = np.random.rand(M)
     ...: 

In [280]: np.allclose(generate_a(b, c, e),generate_a_vectorized(b, c, e))
Out[280]: True

In [281]: %timeit generate_a(b,c,e)
10000 loops, best of 3: 79.4 µs per loop

In [282]: %timeit generate_a_vectorized(b,c,e)
10000 loops, best of 3: 157 µs per loop