例如,对于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).然后我们平均剩下的中间行.
- pd.DataFrame({
- 'value': [10,14],'window_of_last_5_values': [
- np.NaN,np.NaN,'10,7','12,18','8,19','13,9','7,15','18,14'
- ],'values that are counting for average': [
- np.NaN,8',13','result': [
- 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 vs list vs NumPy数组的注释
只需将数据从DataFrame移动到列表,我就可以使用相同的算法获得3.5倍的速度提升.有趣的是,使用NumPy阵列也可以提供几乎相同的速度提升.但是,必须有更好的方法来实现这一目标并实现数量级的提升.
解决方法
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)+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)+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毫秒,比原始大熊猫SciPy方法提高了x60:
- > %timeit rolling_mean_np_jitted_search(data.values)
- 12 ms ± 210 µs per loop (mean ± std. dev. of 7 runs,100 loops each)