首页 > 解决方案 > 为什么某些数据框数学函数花费更多时间?如何加快他们的速度?

问题描述

df1 = pd.DataFrame(data=random_state.randint(10000, size=(3774, 3000)), index=pd.date_range('2010-01-01', '2020-05-01', freq='d'))
print(df1.rolling(window=20).apply(lambda x:x.argmax()))

考虑上面的代码,当我想计算每列的滚动 argmax 时,代码运行得非常慢。

但是当我将 argmax 更改为 max 并运行以下代码时,代码可以在几秒钟内完成:

df1 = pd.DataFrame(data=random_state.randint(10000, size=(3774, 3000)), index=pd.date_range('2010-01-01', '2020-05-01', freq='d'))
# print(df1.rolling(window=20).apply(lambda x:x.argmax()))
print(df1.rolling(window=20).max())

由于 rolling() 对象没有 argmax()、prod() 之类的功能,所以我必须改用 apply(lambda x: x.argmax() / x.prod() ),但这会花费更多时间。

为什么时间相差这么大?如果有任何解决方案可以更快地运行代码?

标签: pythonpandasdataframenumpymath

解决方案


只用numpy>=1.20.0

演示的输入数据:

import pandas as pd
import numpy as np
from numpy.lib.stride_tricks import sliding_window_view

WINDOWSIZE = 3
df = pd.DataFrame(data=10 * np.arange(5*10).reshape((5, 10)))
>>> df
     0    1    2    3    4    5    6    7    8    9
0    0   10   20   30   40   50   60   70   80   90
1  100  110  120  130  140  150  160  170  180  190
2  200  210  220  230  240  250  260  270  280  290
3  300  310  320  330  340  350  360  370  380  390
4  400  410  420  430  440  450  460  470  480  490

用于sliding_window_view在具有给定窗口形状的数组中创建滑动窗口视图:

>>> sliding_window_view(df, (WINDOWSIZE, len(df.columns)))
array([[[[  0,  10,  20,  30,  40,  50,  60,  70,  80,  90],
         [100, 110, 120, 130, 140, 150, 160, 170, 180, 190],
         [200, 210, 220, 230, 240, 250, 260, 270, 280, 290]]],


       [[[100, 110, 120, 130, 140, 150, 160, 170, 180, 190],
         [200, 210, 220, 230, 240, 250, 260, 270, 280, 290],
         [300, 310, 320, 330, 340, 350, 360, 370, 380, 390]]],


       [[[200, 210, 220, 230, 240, 250, 260, 270, 280, 290],
         [300, 310, 320, 330, 340, 350, 360, 370, 380, 390],
         [400, 410, 420, 430, 440, 450, 460, 470, 480, 490]]]])

应用argmax在第三个轴(索引 = 2)并挤压以获得二维数组(数据帧):

>>> sliding_window_view(df, (WINDOWSIZE, len(df.columns))).argmax(axis=2)
array([[[2, 2, 2, 2, 2, 2, 2, 2, 2, 2]],

       [[2, 2, 2, 2, 2, 2, 2, 2, 2, 2]],

       [[2, 2, 2, 2, 2, 2, 2, 2, 2, 2]]])

>>> sliding_window_view(df, (WINDOWSIZE, len(df.columns))).argmax(axis=2).squeeze()
array([[2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
       [2, 2, 2, 2, 2, 2, 2, 2, 2, 2],
       [2, 2, 2, 2, 2, 2, 2, 2, 2, 2]])

最后,将数组转换为 DataFrame:

out = pd.DataFrame(index=df.index, columns=df.columns)
out.iloc[WINDOWSIZE-1:] = sliding_window_view(df, (WINDOWSIZE, len(df.columns))) \
                              .argmax(axis=2).squeeze()
>>> out
   0  1  2  3  4  5  6  7  8  9
2  2  2  2  2  2  2  2  2  2  2
3  2  2  2  2  2  2  2  2  2  2
4  2  2  2  2  2  2  2  2  2  2

表现

WINDOWSIZE = 20
df1 = pd.DataFrame(data=np.random.randint(10000, size=(3774, 3000)), index=pd.date_range('2010-01-01', '2020-05-01', freq='d'))

>>> %timeit sliding_window_view(df1, (WINDOWSIZE, len(df1.columns))).argmax(axis=2).squeeze()
1.43 s ± 5.63 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

推荐阅读