python - 替换 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 ，这是一篇表明递归是向量化的帖子，前提是可以跟踪递归。正如链接的解决方案中所述，目前也不应该很快，但是利用并行处理器，可以尝试这种方法。因此，请将这篇文章作为如何跟踪和向量化递归的指南。这是实现-

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