python的for循环累加_在python中将6 for循环累计和矢量化

编辑3：

最后(我认为)版本，有一点更干净，更快地融入来自max9111's answer的思想。在import numpy as np

from numba import as nb

@nb.njit()

def func1_jit(a, b, c, d):

# Precompute

exp_min = 5 - (a + b + c + d)

exp_max = b

exp = 2. ** np.arange(exp_min, exp_max + 1)

fact_e = np.empty((a + b - 2))

fact_e[0] = 1

for ei in range(1, len(fact_e)):

fact_e[ei] = ei * fact_e[ei - 1]

# Loops

B = 0

for ai in range(0, a):

for bi in range(0, b):

for ci in range(0, c):

for di in range(0, d):

for ei in range(0, ai + bi):

for fi in range(0, ci + di):

B += exp[ei - fi - ai - ci - di + 1 - exp_min] * (ei * ei - 2 * (ei * fi) - 7 * di) * fact_e[ei]

return B

这已经比以前的任何选项都快，但是我们仍然没有利用多个cpu。一种方法是在函数本身内部完成，例如并行化外部循环。这会在创建线程的每个调用上增加一些开销，因此对于较小的输入，实际上会稍慢一点，但对于较大的值，应该会快得多：

^{pr2}$

或者，如果你有很多需要计算函数的点，你也可以在这个级别上并行化。这里a_arr、b_arr、c_arr和{}是要计算函数的值向量：from numba import as nb

@nb.njit(parallel=True)

def func1_arr(a_arr, b_arr, c_arr, d_arr):

B_arr = np.empty((len(a_arr),))

for i in nb.prange(len(B_arr)):

B_arr[i] = func1_jit(a_arr[i], b_arr[i], c_arr[i], d_arr[i])

return B_arr

最佳配置取决于您的输入、使用模式、硬件等，因此您可以结合不同的想法来适合您的情况。在

编辑2：

其实，忘了我之前说过的话。最好的办法是JIT编译算法，但要以更有效的方式进行。首先计算昂贵的部分(我取指数和阶乘)，然后将其传递给编译的loopy函数：import numpy as np

from numba import njit

def func1(a, b, c, d):

exp_min = 5 - (a + b + c + d)

exp_max = b

exp = 2. ** np.arange(exp_min, exp_max + 1)

ee = np.arange(a + b - 2)

fact_e = scipy.special.factorial(ee)

return func1_inner(a, b, c, d, exp_min, exp, fact_e)

@njit()

def func1_inner(a, b, c, d, exp_min, exp, fact_e):

B = 0

for ai in range(0, a):

for bi in range(0, b):

for ci in range(0, c):

for di in range(0, d):

for ei in range(0, ai + bi):

for fi in range(0, ci + di):

B += exp[ei - fi - ai - ci - di + 1 - exp_min] * (ei * ei - 2 * (ei * fi) - 7 * di) * fact_e[ei]

return B

在我的实验中，这是迄今为止最快的选择，并且占用很少的额外内存(只有预计算的值，输入的大小是线性的)。在a, b, c, d = 4, 6, 3, 4

# The original function

%timeit func1_orig(a, b, c, d)

# 2.07 ms ± 33.7 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

# The grid-evaluated function

%timeit func1_grid(a, b, c, d)

# 256 µs ± 25 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

# The precompuation + JIT-compiled function

%timeit func1_jit(a, b, c, d)

# 19.6 µs ± 3.25 µs per loop (mean ± std. dev. of 7 runs, 100000 loops each)

好吧，总有一个选项可以用网格来评估整个事情：import numpy as np

import scipy.special

def func1(a, b, c, d):

ai, bi, ci, di, ei, fi = np.ogrid[:a, :b, :c, :d, :a + b - 2, :c + d - 2]

# Compute

B = (2.) ** (ei - fi - ai - ci - di + 1) * (ei ** 2 - 2 * (ei * fi) - 7 * di) * scipy.special.factorial(ei)

# Mask out of range elements for last two inner loops

m = (ei < ai + bi) & (fi < ci + di)

return np.sum(B * m)

print(func1(4, 6, 3, 4))

# 21769947.844726562

我使用^{}是因为显然^{}由于某种原因不能与数组一起工作。在

显然，随着参数的增加，这种方法的内存开销将迅速增长。代码实际执行的计算比需要的多，因为两个内部循环的迭代次数不同，所以(在这种方法中)必须使用最大的，然后删除不需要的。希望矢量化可以弥补这一点。一个小的IPython基准测试：a, b, c, d = 4, 6, 3, 4

# func1_orig is the original loop-based version

%timeit func1_orig(a, b, c, d)

# 2.9 ms ± 110 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

# func1 here is the vectorized version

%timeit func1(a, b, c, d)

# 210 µs ± 6.34 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

编辑：

注意，前面的方法也不是全有或全无。可以选择仅对部分循环进行栅格求值。例如，最里面的两个循环可以这样矢量化：def func1(a, b, c, d):

B = 0

e = np.arange(a + b - 2).reshape((-1, 1))

f = np.arange(c + d - 2)

for ai in range(0, a):

for bi in range(0, b):

ei = e[:ai + bi]

for ci in range(0, c):

for di in range(0, d):

fi = f[:ci + di]

B += np.sum((2.) ** (ei - fi - ai - ci - di + 1) * (ei ** 2 - 2 * (ei * fi) - 7 * di) * scipy.special.factorial(ei))

return B

这仍然有循环，但它确实避免了额外的计算，并且内存需求要低得多。哪一个最好取决于输入的大小。在我的测试中，对于原始值(4，6，3，4)，这甚至比原始函数慢；而且，对于这种情况，在每个循环上为ei和{}创建新数组要比在预先创建的循环中的一个切片上操作快。但是，如果将输入乘以4(14，24，12，16)，那么这将比原始的(大约x5)快得多，尽管仍然比完全矢量化的(大约x3)慢。另一方面，我可以用这个(大约5分钟)来计算按10(40，60，30，40)缩放的输入值，但由于内存的原因(我没有测试原始函数所需的时间)。使用@numba.jit有点帮助，尽管不是很大(由于阶乘函数，不能使用nopython)。您可以尝试根据输入的大小对更多或更少的循环进行矢量化。在