rolling-statistics.py

@njit_at()
def rolling_statistics(y: ak.Array, builder: ak.ArrayBuilder, u: int):
    y = np.asarray(y)
    result = np.zeros((len(y), 2), dtype=np.float64)
    m = len(y)

    n = u * 2 + 1

    last_var = result[u, 0] = np.var(y[:n])
    last_mean = result[u, 1] = np.mean(y[:n])

    for i in range(1, m):
        # Start at middle of sample (offset by one as we compute var for first sample)
        # e.g. for n = 5
        # 1 2 3 4 5 6
        # ^ ^ $ ^ ^
        # i   j   k
        j = (i + u) % m
        k = (i + n - 1) % m

        end = y[k]
        last_start = y[i - 1]

        mean = last_mean + (end - last_start) / n
        var = last_var + (end - last_start) * (end - mean + last_start - last_mean) / n
                      
                      
        result[j, 0] = var
        result[j, 1] = mean

        last_var = var
        last_mean = mean

    for u in result:
        builder.begin_record()
        builder.field("var").real(u[0])
        builder.field("mean").real(u[1])
        builder.end_record()