Question

数学问题是：

求和中的表达式实际上比上面的表达式复杂得多，但这是一个最小的工作示例，不会使事情复杂化。我已经使用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帖子，其中包括一些示例：

Exterior product in NumPy : Vectorizing six nested loops

Vectorizing triple for loop in Python/Numpy with different array shapes

Python vectorizing nested for loops

我尝试使用从这些有用的帖子中学到的知识，但是经过多次尝试，我一直得出错误的答案。即使对其中一个和进行矢量化处理，也将为真正的问题带来巨大的性能提升，但是和范围不同的事实似乎使我无法接受。有人对如何进行此操作有任何提示吗？

Answer 1

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

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和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

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

编辑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

我之所以使用scipy.special.factorial是因为显然np.factorial由于某种原因不适用于数组。

显然，当您增加参数时，此方法的内存成本会非常快地增长。该代码实际上执行了不必要的计算，因为两个内部循环的迭代次数有所不同，因此（在此方法中）您必须使用最大的迭代，然后删除不需要的代码。希望矢量化可以弥补这一点。一个小的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和fi创建新数组要比对预先创建的数组进行操作要快。但是，如果将输入乘以4（14、24、12、16），则它比原始输入（约x5）要快得多，尽管它仍然比完全矢量化的输入（约x3）要慢。另一方面，我可以用这个值（在5分钟内）将输入的值按10（40、60、30、40）的比例缩放（但在5分钟之内），但是由于内存的原因，不能计算前一个值（我没有测试原始功能需要很长时间）。使用@numba.jit会有所帮助，尽管不是很有用（由于阶乘函数无法使用nopython）。您可以根据输入的大小尝试对更多或更少的循环进行矢量化处理。

Answer 2

这只是对@jdehesa答案的评论。

如果Numba本身不支持该功能，通常建议您自己实现该功能。在进行因式分解的情况下，这并不是一项复杂的任务。

代码

import numpy as np
import numba as nb

@nb.njit()
def factorial(a):
  res=1.
  for i in range(1,a+1):
    res*=i
  return res

@nb.njit()
def func1(a, b, c, d):
    B = 0.

    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)
    for i in range(a + b - 2):
      fact_e[i]=factorial(i)

    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

并行版本

@nb.njit(parallel=True)
def func_p(a_vec,b_vec,c_vec,d_vec):
  res=np.empty(a_vec.shape[0])
  for i in nb.prange(a_vec.shape[0]):
    res[i]=func1(a_vec[i], b_vec[i], c_vec[i], d_vec[i])
  return res

示例

a_vec=np.random.randint(low=2,high=10,size=1000000)
b_vec=np.random.randint(low=2,high=10,size=1000000)
c_vec=np.random.randint(low=2,high=10,size=1000000)
d_vec=np.random.randint(low=2,high=10,size=1000000)

res_2=func_p(a_vec,b_vec,c_vec,d_vec)

单线程版本导致示例中的 5.6µs （首次运行后）。

并行版本将几乎导致另一个Number_of_Cores加速，以计算许多值。请记住，并行版本的编译开销较大（第一次调用的编译开销大于0.5s）。

Answer 3

使用此cartesian_product函数您可以将嵌套循环转换为矩阵，然后可以简单地以矢量化方式计算各自的嵌套sigma：

In [37]: def nested_sig(args):
    ...:     base_prod = cartesian_product(*arrays)
    ...:     second_prod = cartesian_product(base_prod[:,:2].sum(1), base_prod[:,2:].sum(1))
    ...:     total = np.column_stack((base_prod, second_prod))
    ...:     # the items in each row denotes the following variables in order:
    ...:     # ai, bi, ci, di, ei, fi
    ...:     x = total[:, 4] - total[:, 5] - total[:, 0] - total[:, 2] - total[:, 3] + 1
    ...:     y = total[:, 4] - total[:, 5]
    ...:     result = np.power(2, x) * (np.power(total[:, 4], 2) - 2*y - 7*total[:, 3]) * np.math.factorial(total[:,4])
    ...:     return result

Answer 4

我在您的代码中看到了三个改进之处：

range(0,a)是
您在内部循环中做了很多工作
您以随机方式对项求和，存在较大条目丢失精度的风险。

这里有一个版本（可能还不够好），试图改善这一点。

@numba.njit
def func1o(a,b,c,d):
    "2**(ei-fi-ai-ci-di+1)*(ei**2-2*(ei*fi)-7*di)*ei!"                    
    POW=2.;                 SUM=0.;              
    L=[]
    for ai in arange(0.,a+1):
        for bi in range(0,b+1):
            for ci in range(0,c+1):
                for di in range(0,d+1):
                    FACT=1.
                    for ei in arange(0,ai+bi+1):
                        for fi in range(0,ci+di+1):
                            L.append(POW*SUM*FACT)
                            POW /= 2
                            SUM -= 2*ei
                        POW *= 2    
                        SUM += 2*(ei-fi)+1
                        FACT *= ei+1
                    POW /=2
                    SUM -= 7*di
                POW /= 2
        POW /= 2
    A=np.array(L)
    I=np.abs(A).argsort()
    return A[I].sum()

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

4 个答案: