使用Python实现高效的滚动修剪平均数

Question

什么是使用Python计算滚动（移动窗口）修剪均值的最有效方法？

例如，对于一个包含50K行数据和窗口大小为50的数据集，对于每一行，我需要取最后50行，删除前三个和后三个值（窗口大小的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）。然后我们计算剩余中间行的平均值。

生成此示例数据框的代码

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数组也可以获得几乎相同的速度提升。然而，必须有更好的方式来实现这一点，以实现数量级的提升。

- Alex Friedman

---

@roganjosh 你如何在滚动窗口中包含丢弃顶部/底部1%（窗口大小）的值？这可能吗？ - Patrick Artner

我怀疑优化的空间不大，因为计算本身太复杂（例如，不是线性变换）。您也可以尝试使用Cython。 - tif

Answer 1

一个有用的观察是，您不需要在每一步都对所有值进行排序。相反，如果您确保窗口始终有序，您只需要将新值插入到相关位置，并从其原来的位置删除旧值，这两个操作可以使用bisect以 O(log_2(window_size)) 的时间复杂度完成。实际上，这看起来会像这样

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 版本则提高了另一个因素，使我们的时间低于 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友好性。采用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)

- fuglede

---

这里有很多有趣的东西！Numba看起来很棒。我正在尝试复制rolling_mean_np的JIT改进，并且使用窗口大小为50确实获得了可比较的改进 - 但是使用窗口大小为1000却没有。这可能是为什么？ - Alex Friedman

听起来很有道理；在非Numba方法中的向量化可能随着大小的增加变得更加有效。如果您想确切地了解这些方法的区别，并且您还没有使用它，我建议使用line_profiler（这里有一个好指南）；与许多其他分析工具相比，它很轻巧且易于入门。 - fuglede

@Alex Friedman，numba内部的排序函数在许多情况下比numpy版本慢。将排序放在jit编译版本之外可能是有意义的。 - max9111

好问题；就像 decimal.Decimal 一样，对吧？我不知道是否有专门为此目的量身定制的东西，但如果您知道所需的小数位数，您可以通过首先将十进制数组转换为整数之一（即将 [Decimal('1.2'), Decimal('2.5')] 转换为 [12, 25]），计算其修剪平均值，然后再转换回来。这肯定会更慢，但很可能二分搜索仍然是瓶颈。 - fuglede

这个问题的答案（https://dev59.com/sWAf5IYBdhLWcg3wLAGq）提供了一些思路，可以用来得到最小指数。根据您获取数据的位置，您可能能够在上游解决这个问题？ - fuglede

Answer 2

您可以尝试使用 scipy.stats.trim_mean 函数：

from scipy.stats import trim_mean

df['value'].rolling(5).apply(lambda x: trim_mean(x, 0.2))

[输出]

0          NaN
1          NaN
2          NaN
3          NaN
4    10.000000
5    11.000000
6    13.000000
7    13.333333
8    14.000000
9    15.666667

请注意，我必须在您的玩具数据集中使用rolling(5)和proportiontocut=0.2。

对于您的真实数据，您应该使用rolling(50)和trim_mean(x, 0.06)来从滚动窗口中删除前三个和后三个值。

- Chris Adams

---

这是只有我这样觉得，还是实际上并没有得到预期的结果？也就是说，使用 rolling(50) 和 trim_mean(x, 0.05)，第一个非 NaN 值实际上不是 np.mean(sorted(df.value[:50])[3:47])。 - fuglede

有趣！trim_mean保守地切片（将要切片的元素数量向下取整），但是可以调整proportiontocut比例到更高的值来获得所需的数量！我会进行一些测试。 - Alex Friedman

@ChrisA：没错，看起来更好了！ - fuglede

@ChrisA，与我拥有的相比，rolling+trim_mean 的速度惊人地快！有没有办法以某种方式使用 rolling+trimboth 来获取每个窗口的修剪内容？看起来 rolling 无法返回数组，但我希望有一种方法可以解决它。我也会深入研究 fuglede 的答案。 - Alex Friedman

@AlexFriedman 抱歉，我脑海中没有任何方法可以实现那个。不过你应该在这里提出一个单独的问题。肯定会有人有解决方案 :) - Chris Adams

@ChrisA 我想我可以简单地使用rolling.apply函数运行trimboth并将结果保存在某个地方，然后返回均值（或null或其他）。这确实可行，但性能不是很好。继续使用fuglede的手动方法 :) - Alex Friedman

Answer 3

我打赌，每次移动窗口时切片和排序是最慢的部分。不要每次都进行切片，而是制作一个包含50（或5）个值的单独列表。在开始时进行一次排序，然后在添加和删除值（移动窗口）时将新值添加到正确的位置，以保留排序顺序（就像插入排序算法一样）。然后根据该列表中的子集计算修剪均值。您需要一种方法来保持列表与整个集合的关系信息，我认为一个单独的整数变量就足够了。

- Ilia Gilmijarow