python - 使用 Python 的高效滚动修剪均值

Question

用 Python 计算滚动(又名移动窗口)修剪均值的最有效方法是什么？

例如，对于 50K 行的数据集和 50 的窗口大小，对于每一行我需要取最后 50 行，删除顶部和底部的 3 个值(窗口大小的 5%，四舍五入) , 并得到剩余 44 个值的平均值。

目前，我正在对每一行进行切片以获取窗口，对窗口进行排序，然后切片以修剪它。它的工作速度很慢，但必须有更有效的方法。

示例

[10,12,8,13,7,18,19,9,15,14] # data used for example, in real its a 50k lines df

对于 5 的窗口大小。对于每一行，我们查看最后 5 行，对它们进行排序并丢弃 1 个顶部和 1 个底部行(5% 的 5 = 0.25，四舍五入为 1)。然后我们对剩余的中间行进行平均。

将此示例集生成为 DataFrame 的代码

pd.DataFrame({
    'value': [10, 12, 8, 13, 7, 18, 19, 9, 15, 14],
    'window_of_last_5_values': [
        np.NaN, np.NaN, np.NaN, np.NaN, '10,12,8,13,7', '12,8,13,7,18',
        '8,13,7,18,19', '13,7,18,19,9', '7,18,19,9,15', '18,19,9,15,14'
    ],
    'values that are counting for average': [
        np.NaN, np.NaN, np.NaN, np.NaN, '10,12,8', '12,8,13', '8,13,18',
        '13,18,9', '18,9,15', '18,15,14'
    ],
    'result': [
        np.NaN, np.NaN, np.NaN, np.NaN, 10.0, 11.0, 13.0, 13.333333333333334,
        14.0, 15.666666666666666
    ]
})

简单实现的示例代码

window_size = 5
outliers_to_remove = 1

for index in range(window_size - 1, len(df)):
    current_window = df.iloc[index - window_size + 1:index + 1]
    trimmed_mean = current_window.sort_values('value')[
        outliers_to_remove:window_size - outliers_to_remove]['value'].mean()
    # save the result and the window content somewhere

关于 DataFrame 与列表与 NumPy 数组的说明

只需将数据从 DataFrame 移动到列表，我就可以使用相同的算法获得 3.5 倍的速度提升。有趣的是，使用 NumPy 数组也可以提供几乎相同的速度提升。不过，必须有更好的方法来实现这一点并实现数量级的提升。

Answer 1

一个可能会派上用场的观察是，您不需要在每一步都对所有值进行排序。相反，如果您确保窗口始终排序，您需要做的就是在相关位置插入新值，并从原来的位置删除旧值，这两个操作都可以在 O(log_2 (window_size)) 使用 bisect .实际上，这看起来像

def rolling_mean(data):
    x = sorted(data[:49])
    res = np.repeat(np.nan, len(data))
    for i in range(49, len(data)):
        if i != 49:
            del x[bisect.bisect_left(x, data[i - 50])]
        bisect.insort_right(x, data[i])
        res[i] = np.mean(x[3:47])
    return res

现在，在这种情况下，额外的好处比 scipy.stats.trim_mean 所依赖的矢量化所获得的要少，因此特别是，这仍然比@ChrisA 的解决方案，但它是进一步优化性能的有用起点。

> data = pd.Series(np.random.randint(0, 1000, 50000))
> %timeit data.rolling(50).apply(lambda w: trim_mean(w, 0.06))
727 ms ± 34.7 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
> %timeit rolling_mean(data.values)
812 ms ± 42.1 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

值得注意的是，Numba 的抖动在这些情况下通常很有用，但也没有任何好处:

> from numba import jit
> rolling_mean_jit = jit(rolling_mean)
> %timeit rolling_mean_jit(data.values)
1.05 s ± 183 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

以下看似远非最佳的方法优于上述两种其他方法:

def rolling_mean_np(data):
    res = np.repeat(np.nan, len(data))
    for i in range(len(data)-49):
        x = np.sort(data[i:i+50])
        res[i+49] = x[3:47].mean()
    return res

时间:

> %timeit rolling_mean_np(data.values)
564 ms ± 4.44 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

此外，这一次，JIT 编译确实有帮助:

> rolling_mean_np_jit = jit(rolling_mean_np)
> %timeit rolling_mean_np_jit(data.values)
94.9 ms ± 605 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

当我们这样做的时候，让我们快速验证一下这是否确实符合我们的预期:

> np.all(rolling_mean_np_jit(data.values)[49:] == data.rolling(50).apply(lambda w: trim_mean(w, 0.06)).values[49:])
True

事实上，通过稍微帮助排序器，我们可以再挤出一个 2 的因子，将总时间减少到 57 毫秒:

def rolling_mean_np_manual(data):
    x = np.sort(data[:50])
    res = np.repeat(np.nan, len(data))
    for i in range(50, len(data)+1):
        res[i-1] = x[3:47].mean()
        if i != len(data):
            idx_old = np.searchsorted(x, data[i-50])
            x[idx_old] = data[i]
            x.sort()
    return res

> %timeit rolling_mean_np_manual(data.values)
580 ms ± 23 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
> rolling_mean_np_manual_jit = jit(rolling_mean_np_manual)
> %timeit rolling_mean_np_manual_jit(data.values)
57 ms ± 5.89 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
> np.all(rolling_mean_np_manual_jit(data.values)[49:] == data.rolling(50).apply(lambda w: trim_mean(w, 0.06)).values[49:])
True

现在，这个例子中正在进行的“排序”当然只是归结为将新元素放在正确的位置，同时将两者之间的所有内容移动一个。手动执行此操作会使纯 Python 代码变慢，但 jitted 版本获得另一个因子 2，使我们低于 30 毫秒:

def rolling_mean_np_shift(data):
    x = np.sort(data[:50])
    res = np.repeat(np.nan, len(data))
    for i in range(50, len(data)+1):
        res[i-1] = x[3:47].mean()
        if i != len(data):
            idx_old, idx_new = np.searchsorted(x, [data[i-50], data[i]])
            if idx_old < idx_new:
                x[idx_old:idx_new-1] = x[idx_old+1:idx_new]
                x[idx_new-1] = data[i]
            elif idx_new < idx_old:
                x[idx_new+1:idx_old+1] = x[idx_new:idx_old]
                x[idx_new] = data[i]
            else:
                x[idx_new] = data[i]
    return res

> %timeit rolling_mean_np_shift(data.values)
937 ms ± 97.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
> rolling_mean_np_shift_jit = jit(rolling_mean_np_shift)
> %timeit rolling_mean_np_shift_jit(data.values)
26.4 ms ± 693 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
> np.all(rolling_mean_np_shift_jit(data.values)[49:] == data.rolling(50).apply(lambda w: trim_mean(w, 0.06)).values[49:])
True

此时，大部分时间花在 np.searchsorted 上，所以让我们让搜索本身对 JIT 友好。采用the source code for bisect , 我们让

@jit
def binary_search(a, x):
    lo = 0
    hi = 50
    while lo < hi:
        mid = (lo+hi)//2
        if a[mid] < x: lo = mid+1
        else: hi = mid
    return lo

@jit
def rolling_mean_np_jitted_search(data):
    x = np.sort(data[:50])
    res = np.repeat(np.nan, len(data))
    for i in range(50, len(data)+1):
        res[i-1] = x[3:47].mean()
        if i != len(data):
            idx_old = binary_search(x, data[i-50])
            idx_new = binary_search(x, data[i])
            if idx_old < idx_new:
                x[idx_old:idx_new-1] = x[idx_old+1:idx_new]
                x[idx_new-1] = data[i]
            elif idx_new < idx_old:
                x[idx_new+1:idx_old+1] = x[idx_new:idx_old]
                x[idx_new] = data[i]
            else:
                x[idx_new] = data[i]
    return res

这将我们缩短到 12 毫秒，比原始 pandas+SciPy 方法提高了 60 倍:

> %timeit rolling_mean_np_jitted_search(data.values)
12 ms ± 210 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)