Lonewolf
Lonewolf

Reputation: 91

Compute a sequential rolling mean in pandas as array function?

I am trying to calculate a rolling mean on dataframe with NaNs in pandas, but pandas seems to reset the window when it meets a NaN, hears some code as an example...

import numpy as np
from pandas import *

foo = DataFrame(np.arange(0.0,13.0))
foo['1'] = np.arange(13.0,26.0)
foo.ix[4:6,0] = np.nan
foo.ix[4:7,1] = np.nan
bar = rolling_mean(foo, 4)

gives the rolling mean that resets the window after each NaN's, not just skipping out the NaNs

bar = 
       0     1
0    NaN   NaN
1    NaN   NaN
2    NaN   NaN
3    1.5  14.5
4    NaN   NaN
5    NaN   NaN
6    NaN   NaN
7    NaN   NaN
8    NaN   NaN
9    NaN   NaN
10   8.5   NaN
11   9.5  22.5
12  10.5  23.5

I have found an ugly iter/ dropna() work around that gives the right answer

def sparse_rolling_mean(df_data, window):
     ...:     f_data = DataFrame(np.nan,index=df_data.index, columns=df_data.columns)
     ...:     for i in f_data.columns:
     ...:         f_data.ix[:,i] = rolling_mean(df_data.ix[:,i].dropna(),window)
     ...:     return f_data

bar = sparse_rolling_mean(foo,4)

bar
        0     1
0     NaN   NaN
1     NaN   NaN
2     NaN   NaN
3    1.50  14.5
4     NaN   NaN
5     NaN   NaN
6     NaN   NaN
7    3.25   NaN
8    5.00  16.5
9    6.75  18.5
10   8.50  20.5
11   9.50  22.5
12  10.50  23.5

does anyone know if it is possible to do this as an array function ? many thanks in advance.

Upvotes: 2

Views: 1493

Answers (2)

Jeff
Jeff

Reputation: 129018

you can control what get's naned out with the min_periods arg

In [12]:  rolling_mean(foo, 4,min_periods=1)
Out[12]: 
       0     1
0    0.0  13.0
1    0.5  13.5
2    1.0  14.0
3    1.5  14.5
4    2.0  15.0
5    2.5  15.5
6    3.0  16.0
7    7.0   NaN
8    7.5  21.0
9    8.0  21.5
10   8.5  22.0
11   9.5  22.5
12  10.5  23.5

[13 rows x 2 columns]

You can do this if you want results, except when the original was nan

In [27]:  rolling_mean(foo, 4,min_periods=1)[foo.notnull()]
Out[27]: 
       0     1
0    0.0  13.0
1    0.5  13.5
2    1.0  14.0
3    1.5  14.5
4    NaN   NaN
5    NaN   NaN
6    NaN   NaN
7    7.0   NaN
8    7.5  21.0
9    8.0  21.5
10   8.5  22.0
11   9.5  22.5
12  10.5  23.5

[13 rows x 2 columns]

Your expected are a bit odd, as the first 3 rows should have values.

Upvotes: 0

behzad.nouri
behzad.nouri

Reputation: 78011

you may do:

>>> def sparse_rolling_mean(ts, window):
...     return rolling_mean(ts.dropna(), window).reindex_like(ts)
... 
>>> foo.apply(sparse_rolling_mean, args=(4,))
        0     1
0     NaN   NaN
1     NaN   NaN
2     NaN   NaN
3    1.50  14.5
4     NaN   NaN
5     NaN   NaN
6     NaN   NaN
7    3.25   NaN
8    5.00  16.5
9    6.75  18.5
10   8.50  20.5
11   9.50  22.5
12  10.50  23.5

[13 rows x 2 columns]

Upvotes: 2

Related Questions