Matias Andina
Matias Andina

Reputation: 4230

Apply rolling custom function with pandas

There are a few similar questions in this site, but I couldn't find out a solution to my particular question.

I have a dataframe that I want to process with a custom function (the real function has a bit more pre-procesing, but the gist is contained in the toy example fun).

import statsmodels.api as sm
import numpy as np
import pandas as pd
mtcars=pd.DataFrame(sm.datasets.get_rdataset("mtcars", "datasets", cache=True).data)

def fun(col1, col2, w1=10, w2=2):
    return(np.mean(w1 * col1 + w2 * col2))

# This is the behavior I would expect for the full dataset, currently working
mtcars.apply(lambda x: fun(x.cyl, x.mpg), axis=1)

# This was my approach to do the same with a rolling function
mtcars.rolling(3).apply(lambda x: fun(x.cyl, x.mpg))

The rolling version returns this error:

AttributeError: 'Series' object has no attribute 'cyl'

I figured I don't fully understand how rolling works, since adding a print statement to the beginning of my function shows that fun is not getting the full dataset but an unnamed series of 3. What is the approach to apply this rolling function in pandas?

Just in case, I am running

>>> pd.__version__
'1.5.2'

Update

Looks like there is a very similar question here which might partially overlap with what I'm trying to do.

For completeness, here's how I would do this in R with the expected output.

library(dplyr)

fun <- function(col1, col2, w1=10, w2=2){
  return(mean(w1*col1 + w2*col2))
}

mtcars %>% 
  mutate(roll = slider::slide2(.x = cyl,
                               .y = mpg, 
                               .f = fun, 
                               .before = 1, 
                               .after = 1))


                     mpg cyl  disp  hp drat    wt  qsec vs am gear carb     roll
Mazda RX4           21.0   6 160.0 110 3.90 2.620 16.46  0  1    4    4      102
Mazda RX4 Wag       21.0   6 160.0 110 3.90 2.875 17.02  0  1    4    4 96.53333
Datsun 710          22.8   4 108.0  93 3.85 2.320 18.61  1  1    4    1     96.8
Hornet 4 Drive      21.4   6 258.0 110 3.08 3.215 19.44  1  0    3    1 101.9333
Hornet Sportabout   18.7   8 360.0 175 3.15 3.440 17.02  0  0    3    2 105.4667
Valiant             18.1   6 225.0 105 2.76 3.460 20.22  1  0    3    1    107.4
Duster 360          14.3   8 360.0 245 3.21 3.570 15.84  0  0    3    4 97.86667
Merc 240D           24.4   4 146.7  62 3.69 3.190 20.00  1  0    4    2 94.33333
Merc 230            22.8   4 140.8  95 3.92 3.150 22.90  1  0    4    2 90.93333
Merc 280            19.2   6 167.6 123 3.92 3.440 18.30  1  0    4    4     93.2
Merc 280C           17.8   6 167.6 123 3.92 3.440 18.90  1  0    4    4 102.2667
Merc 450SE          16.4   8 275.8 180 3.07 4.070 17.40  0  0    3    3 107.6667
Merc 450SL          17.3   8 275.8 180 3.07 3.730 17.60  0  0    3    3    112.6
Merc 450SLC         15.2   8 275.8 180 3.07 3.780 18.00  0  0    3    3    108.6
Cadillac Fleetwood  10.4   8 472.0 205 2.93 5.250 17.98  0  0    3    4      104
Lincoln Continental 10.4   8 460.0 215 3.00 5.424 17.82  0  0    3    4 103.6667
Chrysler Imperial   14.7   8 440.0 230 3.23 5.345 17.42  0  0    3    4      105
Fiat 128            32.4   4  78.7  66 4.08 2.200 19.47  1  1    4    1      105
Honda Civic         30.4   4  75.7  52 4.93 1.615 18.52  1  1    4    2 104.4667
Toyota Corolla      33.9   4  71.1  65 4.22 1.835 19.90  1  1    4    1     97.2
Toyota Corona       21.5   4 120.1  97 3.70 2.465 20.01  1  0    3    1    100.6
Dodge Challenger    15.5   8 318.0 150 2.76 3.520 16.87  0  0    3    2 101.4667
AMC Javelin         15.2   8 304.0 150 3.15 3.435 17.30  0  0    3    2 109.3333
Camaro Z28          13.3   8 350.0 245 3.73 3.840 15.41  0  0    3    4    111.8
Pontiac Firebird    19.2   8 400.0 175 3.08 3.845 17.05  0  0    3    2 106.5333
Fiat X1-9           27.3   4  79.0  66 4.08 1.935 18.90  1  1    4    1 101.6667
Porsche 914-2       26.0   4 120.3  91 4.43 2.140 16.70  0  1    5    2     95.8
Lotus Europa        30.4   4  95.1 113 3.77 1.513 16.90  1  1    5    2 101.4667
Ford Pantera L      15.8   8 351.0 264 4.22 3.170 14.50  0  1    5    4 103.9333
Ferrari Dino        19.7   6 145.0 175 3.62 2.770 15.50  0  1    5    6      107
Maserati Bora       15.0   8 301.0 335 3.54 3.570 14.60  0  1    5    8     97.4
Volvo 142E          21.4   4 121.0 109 4.11 2.780 18.60  1  1    4    2     96.4

Upvotes: 3

Views: 990

Answers (4)

user5828964
user5828964

Reputation: 312

You can use the below function for rolling apply. It might be slow compared to pandas inbuild rolling in certain situations but has additional functionality.

Function argument win_size, min_periods (similar to pandas and takes only integer input). In addition, after parameter is also used to control to window, it shifts the windows to include after observation.

def roll_apply(df, fn, win_size, min_periods=None, after=None):

    if min_periods is None:
        min_periods = win_size
    else:
        assert min_periods >= 1
    
    if after is None:
        after = 0
    
    before = win_size - 1 - after
    i = np.arange(df.shape[0])
    s = np.maximum(i - before, 0)
    e = np.minimum(i + after, df.shape[0]) + 1
    
    res = [fn(df.iloc[si:ei]) for si, ei in zip(s, e) if (ei-si) >= min_periods]
    idx = df.index[(e-s) >= min_periods]

    types = {type(ri) for ri in res}
    if len(types) != 1:
        return pd.Series(res, index=idx)
    
    t = list(types)[0]
    if t == pd.Series:
        return pd.DataFrame(res, index=idx)
    elif t == pd.DataFrame:
        return pd.concat(res, keys=idx)
    else:
        return pd.Series(res, index=idx)
mtcars['roll'] = roll_apply(mtcars, lambda x: fun(x.cyl, x.mpg), win_size=3, min_periods=1, after=1)
index mpg cyl disp hp drat wt qsec vs am gear carb roll
Mazda RX4 21.0 6 160.0 110 3.9 2.62 16.46 0 1 4 4 102.0
Mazda RX4 Wag 21.0 6 160.0 110 3.9 2.875 17.02 0 1 4 4 96.53333333333335
Datsun 710 22.8 4 108.0 93 3.85 2.32 18.61 1 1 4 1 96.8
Hornet 4 Drive 21.4 6 258.0 110 3.08 3.215 19.44 1 0 3 1 101.93333333333332
Hornet Sportabout 18.7 8 360.0 175 3.15 3.44 17.02 0 0 3 2 105.46666666666665
Valiant 18.1 6 225.0 105 2.76 3.46 20.22 1 0 3 1 107.40000000000002
Duster 360 14.3 8 360.0 245 3.21 3.57 15.84 0 0 3 4 97.86666666666667
Merc 240D 24.4 4 146.7 62 3.69 3.19 20.0 1 0 4 2 94.33333333333333
Merc 230 22.8 4 140.8 95 3.92 3.15 22.9 1 0 4 2 90.93333333333332
Merc 280 19.2 6 167.6 123 3.92 3.44 18.3 1 0 4 4 93.2
Merc 280C 17.8 6 167.6 123 3.92 3.44 18.9 1 0 4 4 102.26666666666667
Merc 450SE 16.4 8 275.8 180 3.07 4.07 17.4 0 0 3 3 107.66666666666667
Merc 450SL 17.3 8 275.8 180 3.07 3.73 17.6 0 0 3 3 112.59999999999998
Merc 450SLC 15.2 8 275.8 180 3.07 3.78 18.0 0 0 3 3 108.60000000000001
Cadillac Fleetwood 10.4 8 472.0 205 2.93 5.25 17.98 0 0 3 4 104.0
Lincoln Continental 10.4 8 460.0 215 3.0 5.424 17.82 0 0 3 4 103.66666666666667
Chrysler Imperial 14.7 8 440.0 230 3.23 5.345 17.42 0 0 3 4 105.0
Fiat 128 32.4 4 78.7 66 4.08 2.2 19.47 1 1 4 1 105.0
Honda Civic 30.4 4 75.7 52 4.93 1.615 18.52 1 1 4 2 104.46666666666665
Toyota Corolla 33.9 4 71.1 65 4.22 1.835 19.9 1 1 4 1 97.2
Toyota Corona 21.5 4 120.1 97 3.7 2.465 20.01 1 0 3 1 100.60000000000001
Dodge Challenger 15.5 8 318.0 150 2.76 3.52 16.87 0 0 3 2 101.46666666666665
AMC Javelin 15.2 8 304.0 150 3.15 3.435 17.3 0 0 3 2 109.33333333333333
Camaro Z28 13.3 8 350.0 245 3.73 3.84 15.41 0 0 3 4 111.8
Pontiac Firebird 19.2 8 400.0 175 3.08 3.845 17.05 0 0 3 2 106.53333333333335
Fiat X1-9 27.3 4 79.0 66 4.08 1.935 18.9 1 1 4 1 101.66666666666667
Porsche 914-2 26.0 4 120.3 91 4.43 2.14 16.7 0 1 5 2 95.8
Lotus Europa 30.4 4 95.1 113 3.77 1.513 16.9 1 1 5 2 101.46666666666665
Ford Pantera L 15.8 8 351.0 264 4.22 3.17 14.5 0 1 5 4 103.93333333333332
Ferrari Dino 19.7 6 145.0 175 3.62 2.77 15.5 0 1 5 6 107.0
Maserati Bora 15.0 8 301.0 335 3.54 3.57 14.6 0 1 5 8 97.39999999999999
Volvo 142E 21.4 4 121.0 109 4.11 2.78 18.6 1 1 4 2 96.4

You can pass more complex function in roll_apply function. Below are few example

roll_apply(mtcars, lambda d: pd.Series({'A': d.sum().sum(), 'B': d.std().std()}), win_size=3, min_periods=1, after=1) # Simple example to illustrate use case

roll_apply(mtcars, lambda d: d, win_size=3, min_periods=3, after=1) # This will return rolling dataframe

Upvotes: 3

I'm not aware of a way to do this calculation easily and efficiently by apply a single function to a pandas dataframe because you're calculating values across multiple rows and columns. An efficient way is to first calculate the column you want to calculate the rolling average for, then calculate the rolling average:

import statsmodels.api as sm
import pandas as pd
mtcars=pd.DataFrame(sm.datasets.get_rdataset("mtcars", "datasets", cache=True).data)

# Create column
def df_fun(df, col1, col2, w1=10, w2=2):
    return w1 * df[col1] + w2 * df[col2]
mtcars['fun_val'] = df_fun(mtcars, 'cyl', 'mpg')

# Calculate rolling average
mtcars['fun_val_r3m'] = mtcars['fun_val'].rolling(3, center=True, min_periods=0).mean()

This gives the correct answer, and is efficient since each step should be optimized for performance. I found that separating the row and column calculations like this is about 10 times faster than the latest approach you proposed and no need to import numpy. If you don't want to keep the intermediate calculation, fun_val, you can overwrite it with the rolling average value, fun_val_r3m.

enter image description here

If you really need to do this in one line with apply, I'm not aware of another way other than what you've done in your latest post. numpy array based approaches may be able to perform better, though less readable.

Upvotes: 2

Matias Andina
Matias Andina

Reputation: 4230

After much searching and fighting against arguments. I found an approach inspired by this answer

def fun(series, w1=10, w2=2):
  col1 = mtcars.loc[series.index, 'cyl']
  col2 = mtcars.loc[series.index, 'mpg']
  return(np.mean(w1 * col1 + w2 * col2))

mtcars['roll'] = mtcars.rolling(3, center=True, min_periods=0)['mpg'] \
                       .apply(fun, raw=False)
mtcars
                      mpg  cyl   disp   hp  ...  am  gear  carb        roll
Mazda RX4            21.0    6  160.0  110  ...   1     4     4  102.000000
Mazda RX4 Wag        21.0    6  160.0  110  ...   1     4     4   96.533333
Datsun 710           22.8    4  108.0   93  ...   1     4     1   96.800000
Hornet 4 Drive       21.4    6  258.0  110  ...   0     3     1  101.933333
Hornet Sportabout    18.7    8  360.0  175  ...   0     3     2  105.466667
Valiant              18.1    6  225.0  105  ...   0     3     1  107.400000
Duster 360           14.3    8  360.0  245  ...   0     3     4   97.866667
Merc 240D            24.4    4  146.7   62  ...   0     4     2   94.333333
Merc 230             22.8    4  140.8   95  ...   0     4     2   90.933333
Merc 280             19.2    6  167.6  123  ...   0     4     4   93.200000
Merc 280C            17.8    6  167.6  123  ...   0     4     4  102.266667
Merc 450SE           16.4    8  275.8  180  ...   0     3     3  107.666667
Merc 450SL           17.3    8  275.8  180  ...   0     3     3  112.600000
Merc 450SLC          15.2    8  275.8  180  ...   0     3     3  108.600000
Cadillac Fleetwood   10.4    8  472.0  205  ...   0     3     4  104.000000
Lincoln Continental  10.4    8  460.0  215  ...   0     3     4  103.666667
Chrysler Imperial    14.7    8  440.0  230  ...   0     3     4  105.000000
Fiat 128             32.4    4   78.7   66  ...   1     4     1  105.000000
Honda Civic          30.4    4   75.7   52  ...   1     4     2  104.466667
Toyota Corolla       33.9    4   71.1   65  ...   1     4     1   97.200000
Toyota Corona        21.5    4  120.1   97  ...   0     3     1  100.600000
Dodge Challenger     15.5    8  318.0  150  ...   0     3     2  101.466667
AMC Javelin          15.2    8  304.0  150  ...   0     3     2  109.333333
Camaro Z28           13.3    8  350.0  245  ...   0     3     4  111.800000
Pontiac Firebird     19.2    8  400.0  175  ...   0     3     2  106.533333
Fiat X1-9            27.3    4   79.0   66  ...   1     4     1  101.666667
Porsche 914-2        26.0    4  120.3   91  ...   1     5     2   95.800000
Lotus Europa         30.4    4   95.1  113  ...   1     5     2  101.466667
Ford Pantera L       15.8    8  351.0  264  ...   1     5     4  103.933333
Ferrari Dino         19.7    6  145.0  175  ...   1     5     6  107.000000
Maserati Bora        15.0    8  301.0  335  ...   1     5     8   97.400000
Volvo 142E           21.4    4  121.0  109  ...   1     4     2   96.400000

[32 rows x 12 columns]


There are several things that are needed for this to perform as I wanted. raw=False will give fun access to the series if only to call .index (False : passes each row or column as a Series to the function.). This is dumb and inefficient, but it works. I needed my window center=True. I also needed the NaN filled with available info, so I set min_periods=0.

There are a few things that I don't like about this approach:

  1. It seems to me that calling mtcars from outside the fun scope is potentially dangerous and might cause bugs.
  2. Multiple indexing with .loc line by line does not scale well and probably has worse performance (doing the rolling more times than needed)

Upvotes: 1

There is no really elegant way to do this. Here is a suggestion:

First install numpy_ext (use pip install numpy_ext or pip install numpy_ext --user).

Second, you'll need to compute your column separatly and concat it to your ariginal dataframe:

import statsmodels.api as sm
import pandas as pd
from numpy_ext import rolling_apply as rolling_apply_ext

import numpy as np

mtcars=pd.DataFrame(sm.datasets.get_rdataset("mtcars", "datasets", cache=True).data).reset_index()
def fun(col1, col2, w1=10, w2=2):
    return(w1 * col1 + w2 * col2)

Col= pd.DataFrame(rolling_apply_ext(fun, 3, mtcars.cyl.values, mtcars.mpg.values)).rename(columns={2:'rolling'})


mtcars.join(Col["rolling"])

to get:

                  index   mpg  cyl   disp   hp  drat     wt   qsec  vs  am  \
0             Mazda RX4  21.0    6  160.0  110  3.90  2.620  16.46   0   1   
1         Mazda RX4 Wag  21.0    6  160.0  110  3.90  2.875  17.02   0   1   
2            Datsun 710  22.8    4  108.0   93  3.85  2.320  18.61   1   1   
3        Hornet 4 Drive  21.4    6  258.0  110  3.08  3.215  19.44   1   0   
4     Hornet Sportabout  18.7    8  360.0  175  3.15  3.440  17.02   0   0   
5               Valiant  18.1    6  225.0  105  2.76  3.460  20.22   1   0   
6            Duster 360  14.3    8  360.0  245  3.21  3.570  15.84   0   0   
7             Merc 240D  24.4    4  146.7   62  3.69  3.190  20.00   1   0   
8              Merc 230  22.8    4  140.8   95  3.92  3.150  22.90   1   0   
9              Merc 280  19.2    6  167.6  123  3.92  3.440  18.30   1   0   
10            Merc 280C  17.8    6  167.6  123  3.92  3.440  18.90   1   0   
11           Merc 450SE  16.4    8  275.8  180  3.07  4.070  17.40   0   0   
12           Merc 450SL  17.3    8  275.8  180  3.07  3.730  17.60   0   0   
13          Merc 450SLC  15.2    8  275.8  180  3.07  3.780  18.00   0   0   
14   Cadillac Fleetwood  10.4    8  472.0  205  2.93  5.250  17.98   0   0   
15  Lincoln Continental  10.4    8  460.0  215  3.00  5.424  17.82   0   0   
16    Chrysler Imperial  14.7    8  440.0  230  3.23  5.345  17.42   0   0   
17             Fiat 128  32.4    4   78.7   66  4.08  2.200  19.47   1   1   
18          Honda Civic  30.4    4   75.7   52  4.93  1.615  18.52   1   1   
19       Toyota Corolla  33.9    4   71.1   65  4.22  1.835  19.90   1   1   
20        Toyota Corona  21.5    4  120.1   97  3.70  2.465  20.01   1   0   
21     Dodge Challenger  15.5    8  318.0  150  2.76  3.520  16.87   0   0   
22          AMC Javelin  15.2    8  304.0  150  3.15  3.435  17.30   0   0   
23           Camaro Z28  13.3    8  350.0  245  3.73  3.840  15.41   0   0   
24     Pontiac Firebird  19.2    8  400.0  175  3.08  3.845  17.05   0   0   
25            Fiat X1-9  27.3    4   79.0   66  4.08  1.935  18.90   1   1   
26        Porsche 914-2  26.0    4  120.3   91  4.43  2.140  16.70   0   1   
27         Lotus Europa  30.4    4   95.1  113  3.77  1.513  16.90   1   1   
28       Ford Pantera L  15.8    8  351.0  264  4.22  3.170  14.50   0   1   
29         Ferrari Dino  19.7    6  145.0  175  3.62  2.770  15.50   0   1   
30        Maserati Bora  15.0    8  301.0  335  3.54  3.570  14.60   0   1   
31           Volvo 142E  21.4    4  121.0  109  4.11  2.780  18.60   1   1   

    gear  carb  rolling  
0      4     4      NaN  
1      4     4      NaN  
2      4     1     85.6  
3      3     1    102.8  
4      3     2    117.4  
5      3     1     96.2  
6      3     4    108.6  
7      4     2     88.8  
8      4     2     85.6  
9      4     4     98.4  
10     4     4     95.6  
11     3     3    112.8  
12     3     3    114.6  
13     3     3    110.4  
14     3     4    100.8  
15     3     4    100.8  
16     3     4    109.4  
17     4     1    104.8  
18     4     2    100.8  
19     4     1    107.8  
20     3     1     83.0  
21     3     2    111.0  
22     3     2    110.4  
23     3     4    106.6  
24     3     2    118.4  
25     4     1     94.6  
26     5     2     92.0  
27     5     2    100.8  
28     5     4    111.6  
29     5     6     99.4  
30     5     8    110.0  
31     4     2     82.8  

Upvotes: 4

Related Questions