Computing a weighted rolling average

[corakingdon]

I am trying to compute a weighted rolling average where each data point has a different weight. (My example is that I have monthly data, and I want each rolling window average to use the number of days in each month as a weight). The following code works, but I wanted to check if there is a better way to do it that doesn't rely on manual iteration of the DataArrayRolling objects.

import xarray as xr
import numpy as np

window_size = 3

# Create one year of monthly data (twelve data points)
da = xr.DataArray(np.arange(0.,12.), coords={"time":np.arange("2022-01", "2023-01", dtype='datetime64[M]')})
month_weights = da.time.dt.daysinmonth # The numbers of days in each month are the desired weights

# Compute a weighted rolling average with manual iteration
da_rolling = da.rolling(time=window_size)
month_weights_rolling = month_weights.rolling(time=window_size)

result = da.copy()
for (da_label, da_window), (_, weights_window),  in zip(da_rolling, month_weights_rolling):
    result.loc[{"time" : da_label}] = da_window.weighted(weights_window.fillna(0)).mean("time")

This is the only way I figured out how to do a rolling average that uses weights unique to each data point (rather than weights that remain the same for each window, which is possible using construct).

I thought something like the following might be available:

result = da.weighted(month_weights).rolling(time = window_size).mean("time")

but DataArrayWeighted does not have a "rolling" method. Is there any other way I can approach this? Thank you in advance for any help!

[max-sixty]

This would be great to have, but unfortunately it's not available in the default methods.

Your approach seems reasonable!

Possibly we could also compute the rolling sum of (values.weights), and divide by the rolling sum of weights...

[dcherian]

The answer to not iterating over rolling and coarsen objects is usually rolling.construct. This will give you a Xarray object with an "expanded" view of the data and then you can run weighted on that.

This tutorial should be useful.

[corakingdon]

perfect, thank you so much!

[corakingdon]

Just to add as a follow-up comment: it appears that the skipna kwarg is only used for rolling averages if you first use construct. (is this an issue with the mean method of DataArrayRolling objects? it doesn't error, but it appears to ignore it)

# add a NAN in the middle of my data:
da[4] = np.NAN

# using `skipna=True` does nothing here
da.rolling(time=3).mean("time", skipna=True)
<xarray.DataArray (time: 12)>
array([nan, nan,  1.,  2., nan, nan, nan,  6.,  7.,  8.,  9., 10.])
Coordinates:
  * time     (time) datetime64[ns] 2022-01-01 2022-02-01 ... 2022-12-01

# same as skipna=False
da.rolling(time=3).mean("time", skipna=False)
<xarray.DataArray (time: 12)>
array([nan, nan,  1.,  2., nan, nan, nan,  6.,  7.,  8.,  9., 10.])
Coordinates:
  * time     (time) datetime64[ns] 2022-01-01 2022-02-01 ... 2022-12-01

# whereas if you call construct first, then the skipna kwarg works
da.rolling(time=3).construct("window").mean("window", skipna=True)
<xarray.DataArray (time: 12)>
array([ 0. ,  0.5,  1. ,  2. ,  2.5,  4. ,  5.5,  6. ,  7. ,  8. ,  9. ,
       10. ])
Coordinates:
  * time     (time) datetime64[ns] 2022-01-01 2022-02-01 ... 2022-12-01

# skipna=False
da.rolling(time=3).construct("window").mean("window", skipna=False)
<xarray.DataArray (time: 12)>
array([nan, nan,  1.,  2., nan, nan, nan,  6.,  7.,  8.,  9., 10.])
Coordinates:
  * time     (time) datetime64[ns] 2022-01-01 2022-02-01 ... 2022-12-01

[dcherian]

That looks like a bug. Can you open an issue please?

Note rolling has a min_count kwarg that also affects this behaviour.

[corakingdon]

sure, will start a new issue!

[corakingdon]

#6772