Pandas groupby mean by specifying minimum number of non-NA values

Viewed 263

I have a dataframe in the following format:

               A     B     C     D
2020-11-18  64.0  74.0  34.0  57.0
2020-11-20   NaN  71.0   NaN  58.0
2020-11-23   NaN  11.0   NaN   NaN
2020-11-25  69.0   NaN   NaN   0.0
2020-11-27   NaN  37.0  19.0   NaN
2020-11-29  63.0   NaN   NaN  85.0
2020-12-03   NaN  73.0   NaN  49.0
2020-12-10   NaN   NaN  32.0   NaN
2020-12-22  52.0  90.0  33.0  24.0
2020-12-23   NaN  96.0   NaN   NaN
2020-12-28  78.0   NaN   NaN  68.0
2020-12-29  17.0  70.0   NaN  16.0
2021-01-03  51.0  43.0   NaN  66.0

I want to compute the mean values in each column for each month. I am aware that if it was a rolling window mean computation (e.g. df.rolling(5).mean()) then I would have an optional parameter in rolling() allowing me to specify a minimum number of required valid (non-NA) values to compute the mean; if there are at least as many non-NA values as specified, the mean value is returned, otherwise a NaN value is returned.

I am looking for the same functionality when using groupby / resample. For instance, using resample, I can get all the monthly means like so:

df.resample('1M').mean()

but I would like to only calculate the mean if the corresponding month has at least 3 valid values. In that case, column C for November 2020 shouldn't have a numerical output since there are only two non-NA cells there. What would be the easiest way to go about this?

2 Answers

You can use groupby().agg() to find mean and num_nan:

tmp = df.groupby(pd.Grouper(freq='M')).agg(['mean','count']).swaplevel(0,1,axis=1)

tmp['mean'].where(tmp['count']>=3)

Output:

                    A      B   C      D
2020-11-30  65.333333  48.25 NaN  50.00
2020-12-31  49.000000  82.25 NaN  39.25
2021-01-31        NaN    NaN NaN    NaN

I would write a custom function with the apply method. I showed it with a sample series. If something is not clear just ask.

def naCheck(series):
    series = series.dropna()
    if len(series) < 3:
        return np.nan
    else:
        return series.mean()

index = pd.date_range('1/1/2000', periods=9, freq='T')
series = pd.Series(range(9), index=index)
series = series.resample('3T').apply(naCheck)
print(series)
Related