在python中向量化6个for循环累积和

数学问题是：

总和中的表达式实际上比上面的表达式复杂得多，但这是一个最小的工作示例，以免事情过于复杂。我使用 6 个嵌套 for 循环在 Python 中编写了此代码，正如预期的那样，即使在 Numba、Cython 和朋友的帮助下，它的性能也非常糟糕（真正的形式性能很差，需要评估数百万次）。这里是使用嵌套 for 循环和累积和编写的：

import numpy as np

def func1(a,b,c,d):
    '''
    Minimal working example of multiple summation
    '''
    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 += (2)**(ei-fi-ai-ci-di+1)*(ei**2-2*(ei*fi)-7*di)*np.math.factorial(ei)


    return a, b, c, d, B

该表达式由 4 个数字作为输入来控制func1(4,6,3,4)的输出B是 21769947.844726562。

我四处寻找这方面的帮助，发现了几个 Stack 帖子，其中一些例子是：

NumPy 中的外部积：矢量化六个嵌套循环 https://stackoverflow.com/questions/41808217/exterior-product-in-numpy-vectorizing-six-nested-loops

在 Python/Numpy 中使用不同的数组形状向量化三重 for 循环 https://stackoverflow.com/questions/32755768/vectorizing-triple-for-loop-in-python-numpy-with-different-array-shapes

Python 向量化嵌套 for 循环 https://stackoverflow.com/questions/39667089/python-vectorizing-nested-for-loops

我尝试使用从这些有用的帖子中学到的知识，但经过多次尝试，我不断得到错误的答案。即使对内部总和之一进行矢量化也会为实际问题带来巨大的性能增益，但总和范围不同的事实似乎让我感到困惑。有人对如何取得进展有任何建议吗？

EDIT 3:

最终（我认为）版本，更干净，更快地融入了来自max9111 的回答 https://stackoverflow.com/a/53320252.

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。一种方法是在函数本身内部实现，例如并行化外循环。这会在每次调用创建线程时增加一些开销，因此对于较小的输入实际上会慢一些，但对于较大的值应该会明显更快：

import numpy as np
from numba import as nb

@nb.njit(parallel=True)
def func1_par(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 = np.empty((a,))
    for ai in nb.prange(0, a):
        Bi = 0
        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):
                            Bi += exp[ei - fi - ai - ci - di + 1 - exp_min] * (ei * ei - 2 * (ei * fi) - 7 * di) * fact_e[ei]
        B[ai] = Bi
    return np.sum(B)

或者，如果您想要评估函数的许多点，您也可以在该级别进行并行化。这里a_arr, b_arr, c_arr and d_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

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

EDIT 2:

其实，忘记我之前说的话了。最好的办法是以更有效的方式 JIT 编译算法。首先计算昂贵的部分（我采用指数和阶乘），然后将其传递给编译后的循环函数：

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

I used scipy.special.factorial https://docs.scipy.org/doc/scipy/reference/generated/scipy.special.factorial.html因为显然np.factorial https://docs.scipy.org/doc/numpy/reference/generated/numpy.factorial.html由于某种原因不适用于数组。

显然，这样的内存成本会增加very增加参数时速度会加快。该代码实际上执行了不必要的计算，因为两个内部循环的迭代次数不同，因此（在此方法中）您必须使用最大的计算，然后删除不需要的计算。希望矢量化能够弥补这一点。一个小型 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)

EDIT:

请注意，之前的方法也不是全有或全无的事情。您可以选择仅对某些循环进行网格评估。例如，两个最里面的循环可以像这样矢量化：

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 and fi在每个循环上进行操作比在预先创建的循环的一部分上进行操作要快。然而，如果将输入乘以 4 (14, 24, 12, 16)，那么这比原始值（大约 x5）快得多，尽管仍然比完全矢量化的（大约 x3）慢。另一方面，我可以使用这个（大约 5 分钟）计算输入缩放十倍（40,60,30,40）的值，但由于内存的原因不能使用前一个（我没有测试如何原始功能需要很长时间）。使用@numba.jit有一点帮助，虽然不是很大（不能使用nopython由于阶乘函数）。您可以根据输入的大小尝试对更多或更少的循环进行矢量化。