Pandas rolling calculation based on a mask condition?

user_5123111 0
Viewed 290

I have a dataframe in the following format:

>>> df = pd.DataFrame(np.random.randint(0,100,size=(100, 4)), columns=list('ABCD'))
     A   B   C   D
0   58   3  25  94
1   54   4  99  85
2   87  98  74  36
3   63  68  79  51
4   25   7  42  12
..  ..  ..  ..  ..
95  50  19  90  99
96  68   1  52  98
97  96   2  98  21
98  90  25   5  23
99  56  93  48  19

I want to recompute those values that meet a certain criterium. There can be many criteria but for the sake of simplicity, let's just take a simple one: if a cell value is below 10% of the dataframe mean:

>>> mask = df<0.1*df.mean().mean()
        A      B      C      D
0   False   True  False  False
1   False   True  False  False
2   False  False  False  False
3   False  False  False  False
4   False  False  False  False
..    ...    ...    ...    ...
95  False  False  False  False
96  False   True  False  False
97  False   True  False  False
98  False  False   True  False
99  False  False  False  False

For each cell that meets this condition, I want to create a rolling window of size 10, compute the mean over that rolling window, and replace the cell's value with that mean. For all the rest of the cells in the original df dataframe, original values should be kept. Therefore we can think of this intuitively as: if we see a value that is too low, let's replace it by the average of the last 10 (in fact only 9, due to the way rolling window is defined in pandas) values.

Now I am a bit how to move forward from here. df[mask] will show me the cells that meet the condition but I am unable to apply a .rolling() on it since df[mask] has NaNs in all other cells that do not meet the condition.

Note that I prefer high performance as in reality both the criteria and the rolling window calculation will be more complicated, and the dataframe is much larger. Therefore I could do a simple for loop type of lookup but that doesn't suit my needs for this use case.

pandas python
3 Answers

If I understand the question correctly:

mask = df < 0.1 * df.mean().mean()
df[mask] = df.rolling(10, 0).mean()[mask]
user_7581740

Depending how you want to treat NaN, but one approach is simply: df[mask].fillna(0).rolling(10).mean()

user_95013640

I managed to do it by turning the problem around: instead of focusing on figuring out how to only perform the rolling window calculation for those cells that meet the criterium, I first calculated the rolling mean for the entire dataframe:

>>> means = df.rolling(10).mean()
       A     B     C     D
0    NaN   NaN   NaN   NaN
1    NaN   NaN   NaN   NaN
2    NaN   NaN   NaN   NaN
3    NaN   NaN   NaN   NaN
4    NaN   NaN   NaN   NaN
..   ...   ...   ...   ...
95  46.3  35.2  44.8  41.5
96  46.1  40.0  42.0  48.0
97  49.6  45.6  42.8  52.6
98  46.7  39.5  40.3  49.9
99  43.6  42.5  35.9  46.7

Then, I simply took df[mask] where mask = df<0.1*df.mean()*mean() and assigned means to it like so:

df[mask] = means

This keeps the original df cell values where mask is False but will replace them with the corresponding cell values from means if mask is True.

user_51231110
Related