慕仙森
编辑 3：最终（我认为）版本，从max9111's answer 中结合了一些想法，更简洁、更快速。import numpy as npfrom 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 npfrom 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 编译，但要以更有效的方式。首先计算昂贵的部分（我采用指数和阶乘），然后将其传递给已编译的循环函数：import numpy as npfrom numba import njitdef 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 npimport scipy.specialdef 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 分钟内）计算按十（40、60、30、40）缩放的输入值，但由于内存的原因不能用前一个（我没有测试如何使用原始函数需要很长时间）。使用@numba.jit有一点帮助，虽然不是很大（nopython由于阶乘函数而无法使用）。您可以根据输入的大小尝试对更多或更少的循环进行矢量化。
慕丝7291255
如果 Numba 本身不支持某个功能，通常建议您自己实现它。在分解的情况下，这不是一项复杂的任务。代码import numpy as npimport 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.5 秒）。