在python中向量化6 for循環累積和

Question

數學問題是：

總和中的表達式實際上比上面的表達式復雜得多，但這是一個最小的工作示例，不會使事情過於復雜。 我用Python編寫了6個嵌套for循環，並且正如預期的那樣表現非常糟糕（真正的表單執行得很糟糕，需要評估數百萬次），即使在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中的外部產品：矢量化六個嵌套循環

在Python / Numpy中使用不同的數組形狀矢量化三重循環

Python矢量化嵌套for循環

我試圖使用我從這些有用的帖子中學到的東西，但經過多次嘗試，我一直得出錯誤的答案。 即使對其中一個內部總和進行矢量化也會為真正的問題帶來巨大的性能提升，但總和范圍不同的事實似乎讓我失望了。 有沒有人有關於如何推進這個的任何提示？

Answer 1

編輯3：

最終（我認為）版本，更清晰，更快速地結合了max9111的答案 。

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編譯算法，但是以更有效的方式。 首先計算昂貴的部分（我采用指數和階乘），然后將其傳遞給編譯的循環函數：

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不會因為某些原因數組。

Obivously，因為你增加了參數的這種內存的成本會增長非常快。 代碼實際上執行的計算比必要的多，因為兩個內部循環具有不同的迭代次數，因此（在此方法中）您必須使用最大的，然后刪除您不需要的。 希望是矢量化將彌補這一點。 一個小的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）慢。 另一方面，我可以計算輸入的值，用十（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.5秒）。

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()