如何使用 Numba 有效地加快簡單的移動平均線計算

Question

我正在嘗試使用 Numba 來加速股票市場分析的一些簡單迭代函數。 我對 Pandas 或 Numpy 不感興趣，我只是想了解無蟒蛇（@njit）function 的方法是什么。

這是一個簡單的移動平均線 function：

def sma_plain(src,p):
    win = []
    res = []
    for n in src:
        win.append(n)
        if len(win) > p:
            win = win[1:]
        res.append(sum(win)/len(win))
    return res

我可以立即告訴我無法使用 Numba 設置空的 win[] 和 res[] 列表。 我嘗試使用 Numba 的 List()，但在 function 中無法將其初始化為空。 我嘗試從復制 output 的 src 開始，並使用一片 src 實例化 window（通過 window，我的意思是一組將被編譯的值）。 同樣在我進行編譯的嘗試之一中，與我原來的 function 相比，timeit 產生的結果更慢。 很可能是因為 sum() 不能使用。 我在文檔中看到的唯一示例是使用 gpuvectorization 和我不理解的語法。 我對此不感興趣，但是，我只是想了解 Numba for nopython 的過程。

Also I made the function where it uses a smaller window for the initial-values, ideally I would like to write None or null values to the list when the window hasn't matured, but it seems that Numba does not allow for null values. 沒關系，我想我可以在這種情況下簡單地使用較短的列表並在進行計算時跟蹤偏移量，但如果能夠跟蹤 null 值，那就太好了。

這是我最后一次嘗試，但它並沒有真正有用，因為它不能編譯。

@njit
def sma(src, p):
    res = src.copy()
    i = 0
    length = len(src)
    while i < length:
        win = src[max(0, i+1-p) : i+1]
        win_length = len(win)
        s = 0
        for n in win:
            s += n
        s /= win_length
        res[i] = s
    return res

編輯：我有一個 function 可以編譯並且現在似乎允許 None 值。 我不知道為什么它之前給了我錯誤。 所以現在我有以下類似的function：

@njit
def sma(src, p):
    slices = [src[i-p:i] for i in range(p,len(src)+1)]
    res = []
    for slc in slices:
        s = 0.0
        for i in range(p):
            s += slc[i]
        res.append(s/p)
    res = [None]*(p-1) + res
    return res

def sma_plain(src,p):
    win = []
    res = []
    for n in src:
        win.append(n)
        if len(win) > p:
            win = win[1:]
        if len(win) == p:
            res.append(sum(win)/len(win))
        else:
            res.append(None)
    return res

但是對於這個在一些股票數據時間上進行 1000 次迭代的情況，它報告 numba function 為 39 秒，Python ZC1C425268E68385D1AB5074F14ZA 為 7 秒。 現在我想知道 Numba 是否還在工作，因為它不會像以前那樣拋出錯誤。 編譯函數是否存在緩存問題，導致它使用過時的版本或其他東西？

Answer 1

只要你不想用Numpy，我就不推薦這個；-)：

def sma_numpy_acc(a, p):
    m = np.cumsum(a) / p
    m[p:] -= m[:-p]
    m[:p-1] = np.nan
    return m    

請注意，我使用的是 NaN 而不是 None，因為數組可以具有同質類型。

與原始功能相比的時序：

%timeit sma(a, p)
88.6 ms ± 1.58 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
%timeit sma_plain(a, p)
2.18 s ± 65.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
%timeit sma_numpy_acc(a, p)
3.95 ms ± 56.3 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

通過jitting function可以稍微提高速度：

@nb.njit
def sma_numpy_acc_jit(a, p):
    m = np.cumsum(a) / p
    m[p:] = m[p:] - m[:-p]        # Odd behavior of -= in Numba
    m[:p - 1] = np.nan
    return m

%timeit sma_numpy_acc_jit(a, p)
3 ms ± 66.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

使用累積和的相同想法，仍然使用 Numpy arrays，但不使用 Numpy 函數：

@nb.njit
def sma_jit_acc(a, p):
    acc = np.empty_like(a)
    acc[0] = a[0]
    n = len(a)
    for i in range(1, n):
        acc[i] = acc[i-1] + a[i]
    for i in range(n-1, p-1, -1):
        acc[i] = (acc[i] - acc[i-p]) / p
    acc[p-1] /= p
    for i in range(p-1):
        acc[i] = np.nan
    return acc

時序類似於純Numpy function。

%timeit sma_jit_acc(a, p)
3.69 ms ± 119 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

使用列表的相同方法。 Numpy 沒有蹤跡：

@nb.njit
def sma_jit_acc_lists(a, p):
    n = len(a)
    acc = [math.nan] * n
    acc[0] = a[0]
    for i in range(1, n):
        acc[i] = acc[i-1] + a[i]
    for i in range(n-1, p-1, -1):
        acc[i] = (acc[i] - acc[i-p]) / p
    acc[p-1] /= p
    return acc

使用列表會降低時間：

%timeit sma_jit_acc_lists(a, p)
24.2 ms ± 1.47 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)