Why some dataframe mathmatical functions cost more time? How to speed up them?

user_9177008 0
Viewed 67
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()))

Consider the above code, when I want to compute the rolling argmax of each column, the code run very slow.

But when I change argmax to max, and run the following code, the code could finish in seconds:

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())

Since rolling() object has no function like argmax(), prod(), so I have to use apply(lambda x: x.argmax() / x.prod() ) instead, but it would cost more more more time.

Why the time differs so much? If there any solutions to run the code faster?

pandas dataframe math numpy python
1 Answers

With numpy>=1.20.0 only.

Input data for demo:

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

Use sliding_window_view to create a sliding window view into the array with the given window shape:

>>> 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]]]])

Apply argmax over the third axis (index=2) and squeeze to get a 2D-array (dataframe like):

>>> 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]])

Finally, convert array to 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

Performance

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)
user_152399510
Related