如何使用 Numba 有效地加快简单的移动平均线计算

Question

I am trying to use Numba to speed up some simple iterative functions for stock market analysis. I'm not interested in Pandas or Numpy for this, I'm just trying to understand what the approach would be for a no-python (@njit) function.

Here is a simple moving average 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

I can tell immediately that I can't set the empty win[] and res[] lists with Numba. I tried to use Numba's List(), but that can't be initialized as empty within the function. I tried starting by copying src for the output, and using a slice of src to instantiate the window (by window I mean a set of values that will be summed), but my code won't compile. Also in one of my attempts that did compile, timeit produced slower results compared to my original function. Most likely because sum() can't be used. The only example I could see in the docs was using gpuvectorization and with a syntax I didn't understand. I'm not interested in that, yet, I simply want to understand the process with 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. This is ok, I think I could simply use a shorter list in this scenario and keep track of an offset when making calculations, but it would be great if there were an ability to track null values.

This was my last attempt, but it's not really useful as it doesn't compile.

@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

Edit: I have a function that compiles and seems to be allowing for None values now. I don't know why it was giving me errors before. So now I have the following similar 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

But for this with 1000 iterations over some stock data timeit reports 39 seconds for the numba function, and 7 seconds for the Python function. Now I'm wondering if Numba is even working any longer as it isn't throwing errors for everything like it was. Are there caching issues with the compiled functions that cause it to use outdated versions or something?

Answer 1

As long as you don't want to use Numpy, I won't recommend this one;-):

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

Note I'm using NaN instead of None, son that the array can have homogeneous types.

Timing compared to the original functions:

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

Speed can be slightly increased by jitting the 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)

The same idea of using a cumulative sum, still using Numpy arrays, but not Numpy functions:

@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

Timing is similar to the pure Numpy function.

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

The same method using lists. No trace of 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

Timing is degraded by the use of lists:

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