Python 的高效滚动截尾平均值
Efficient rolling trimmed mean with Python
计算滚动（又名移动 window）的最有效方法是什么trim用 Python 计算平均值？
例如，对于包含 50K 行且 window 大小为 50 的数据集，对于每一行，我需要取最后 50 行，删除顶部和底部的 3 个值（5% window 大小，四舍五入），并取剩余 44 个值的平均值。
目前，我正在对每一行进行切片以获得 window，对 window 进行排序，然后切片为 trim。它的工作速度很慢，但必须有更有效的方法。
例子
[10,12,8,13,7,18,19,9,15,14] # data used for example, in real its a 50k lines df
对于大小为 5 的 window。对于每一行，我们查看最后 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 数组也可以提供几乎相同的速度提升。尽管如此，必须有更好的方法来实现这一点并实现数量级的提升。
您可以尝试使用 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) 从滚动 window.
中删除顶部和底部的 3 个值
一个可能会派上用场的观察结果是您不需要在每一步都对所有值进行排序。相反，如果您确保 window 始终排序，您需要做的就是在相关位置插入新值，并从原来的位置删除旧值，这两者都是可以完成的操作在 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
现在，这个例子中发生的 "sorting" 当然只是归结为将新元素放在正确的位置，同时将两者之间的所有内容移动一个。手动执行此操作会使纯 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)
我敢打赌 window 的每一步都进行切片和排序是最慢的部分。不要每次都切片，而是制作一个包含 50（或 5）个值的单独列表。在开始时排序一次，然后在添加和删除值时（移动 window）在正确的位置添加新值以保持排序顺序（很像插入排序算法）。然后根据该列表中的值子集计算修剪平均值。
您将需要一种方法来保存有关列表相对于整个集合的位置的信息，我认为单个 int 变量就足够了。
计算滚动（又名移动 window）的最有效方法是什么trim用 Python 计算平均值？
例如，对于包含 50K 行且 window 大小为 50 的数据集，对于每一行，我需要取最后 50 行，删除顶部和底部的 3 个值（5% window 大小，四舍五入），并取剩余 44 个值的平均值。
目前，我正在对每一行进行切片以获得 window，对 window 进行排序，然后切片为 trim。它的工作速度很慢，但必须有更有效的方法。
例子
[10,12,8,13,7,18,19,9,15,14] # data used for example, in real its a 50k lines df
生成此示例的代码设置为 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 数组也可以提供几乎相同的速度提升。尽管如此，必须有更好的方法来实现这一点并实现数量级的提升。
您可以尝试使用 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) 从滚动 window.
一个可能会派上用场的观察结果是您不需要在每一步都对所有值进行排序。相反，如果您确保 window 始终排序，您需要做的就是在相关位置插入新值，并从原来的位置删除旧值，这两者都是可以完成的操作在 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
现在，这个例子中发生的 "sorting" 当然只是归结为将新元素放在正确的位置，同时将两者之间的所有内容移动一个。手动执行此操作会使纯 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)
我敢打赌 window 的每一步都进行切片和排序是最慢的部分。不要每次都切片，而是制作一个包含 50（或 5）个值的单独列表。在开始时排序一次，然后在添加和删除值时（移动 window）在正确的位置添加新值以保持排序顺序（很像插入排序算法）。然后根据该列表中的值子集计算修剪平均值。 您将需要一种方法来保存有关列表相对于整个集合的位置的信息，我认为单个 int 变量就足够了。