Rolling Median in a Pandas Pivot Table

user_9571463 0
Viewed 889

I am trying to calculate a rolling median as an aggregated function on a pandas dataframe. Here is some sample data:

import pandas as pd
import numpy as np

d = {'date': ['2020-01-01','2020-02-01','2020-03-01','2020-01-01','2020-02-01','2020-02-01','2020-03-01','2020-02-01','2020-03-01','2020-03-01','2020-03-01','2020-03-01','2020-03-01'],
     'count': [1,1,1,2,2,3,3,3,4,3,3,3,1], 
     'type': ['type1','type2','type3','type1','type3','type1','type2','type2','type2','type3','type1','type2','type1'],
     'salary':[1000,2000,3000,10000,15000,30000,100000,50000,25000,10000,25000,30000,40000]}
df: pd.DataFrame = pd.DataFrame(data=d)

df_pvt: pd.DataFrame = df.pivot_table(index='date',
                                      columns='type',
                                      aggfunc={'salary': np.median})
df_pvt.head(5)

I would like to perform a rolling median on the salaries using pandas rolling(2).median() function.

How can I go about inserting this type of window function into the aggregate function for a pivot table?

My goal is to aggregate a large amount of numeric data by date and take the rolling median of variable lengths and report that in my resulting pivot table. I am not entirely sure how to insert this function into aggfunc or the like.

The expected output orders by the date in ascending order and takes all observations associated with both months and finds the median.

For type1 we have:


date    count   type    salary
0   2020-01-01  1   type1   1000
3   2020-01-01  2   type1   10000
5   2020-02-01  3   type1   30000
10  2020-03-01  3   type1   25000
12  2020-03-01  1   type1   40000

Thus, for type1 the expected output with rolling(2) would be:


             salary
type         type1  
date            
2020-01-01  NaN 
2020-02-01  10000.0
2020-03-01  30000.0

The logic follows that for the first 2 month rolling window we would have data points 1000,10000 and 30000 and produce a median of 10000.

For 2020-03-01, the rolling 2 would include 30000, 25000, 40000 so the median result should be 30000.

pandas median dataframe numpy python
1 Answers

Not sure it can be done directly with the parameter aggfunc. so a work around could be to create the a double of the data with a date column shifted of a month. Note that this method is not really scalable to bigger rolling window. it can but you may end up with too much data.

# first convert to datetime
df['date'] = pd.to_datetime(df['date'])

# append the data shifted of a month to df and perform the pivot_table
res = (
    df
    .append(df.assign(date=lambda x: x['date']+pd.DateOffset(months=1)))
    .pivot_table(index='date',columns='type',
                 aggfunc={'salary': np.median})
    .reindex(df['date'].unique()) # to avoid an extra month
)

print(res)
             salary                  
type          type1    type2    type3
date                                 
2020-01-01   5500.0      NaN      NaN
2020-02-01  10000.0  26000.0  15000.0
2020-03-01  30000.0  30000.0  10000.0

for the first date if you want to get nan as a rolling window would do, then you can do res.loc[res.index.min()] = np.nan after

user_92747320
Related