How do I reduce a third dimension when doing a 2D rolling mean?

[tbloch1]

I feel like I'm missing something super obvious here, but, for the life of me, I can't figure out how to pass a spatial rolling window over data and average the data in that window across time. Hopefully the examples below will make what I'm trying to do clear.

Given what I've tried, I think the issue might be that I just don't understand how .rolling() works. My assumption was that it would work basically like .sel() for the coords of each window. So in my case, the .rolling() (essentially) return a bunch of DataArrays, with the subset of spatial coords, but all the time coords. Then when you do .mean() it would take all that data and reduce it down to a single mean value, and then it would construct a new dataarray with a shape equivalent to the dimensions being rolled over, and with whatever dimensions were remaining after the .mean() (depending on which dims you performed the mean over etc.).

Below are the main things I've tried so far. I want to get an output that is shaped (100,100), having averaged over the time dimension.

# Create 10 timesteps for a 100x100 image
shp = (100, 100, 10)
image = np.arange(np.prod(shp)).reshape(shp)
 
# Create an xarray dataset from the image
ds = xr.Dataset({'image': (['y', 'x', 't'], image)}, coords={'y': np.arange(100), 'x': np.arange(100), 't': np.arange(10)})
 
# Define the window size and stride
cube_size = 10
 
# Method 1
# Normal rolling mean, assuming it will reduce the final dimension.
r_mean = ds['image'].rolling(y=cube_size, x=cube_size, center=False).mean()
print('Method 1', r_mean.shape)
 
# Method 2
# Maybe I need to specify the dims for the mean?
r_mean = ds['image'].rolling(y=cube_size, x=cube_size, center=False).mean(dim=('y', 'x', 't'))
print('Method 2', r_mean.shape)
 
# Method 3
# Okay, so I need to specify the time dimension for the reduction?
r_mean = ds['image'].rolling(y=cube_size, x=cube_size, t=10, center=False).mean()
print('Method 3', r_mean.shape)
# Maybe I just need to take the final time slice and that will represent the average over all the time data?
print('Number of NaNs should be 0, but there are', np.isnan(r_mean[:,:,9].values).sum())
 
# Method 4
# Maybe I need to use the reduce function instead.
r_mean = ds['image'].rolling(y=cube_size, x=cube_size, center=False).reduce(np.mean)
print('Method 4', r_mean.shape)

And the outputs are:

Method 1 (100, 100, 10)
Method 2 (100, 100, 10)
<ipython-input-38-43c65d803fc2>:18: DeprecationWarning: Reductions are applied along the rolling dimension(s) '['y', 'x']'. Passing the 'dim' kwarg to reduction operations has no effect.
  r_mean = ds['image'].rolling(y=cube_size, x=cube_size, center=False).mean(dim=('y', 'x', 't'))
Method 3 (100, 100, 10)
Number of NaNs should be 0, but there are 1719
Method 4 (100, 100, 10)

[tbloch1]

Okay, I think I found a method to do it. I'm leaving this open because I still feel like there must be a nicer way to do it (or so someone can explain what I don't understand about how rolling() works...).

The logic is that we use the construct() method to extract the data we want, and then do the average over the dimensions we want afterwards. This definitely still seems like a clunky approach to me, so there's probably a better way.

Demo:

# Create dummy data
shp = (100, 100, 10)
images = np.arange(np.prod(shp).reshape(shp)
# Create dataset
ds = xr.Dataset(
    {'images': (['y', 'x', 't'], images)},
    coords = {
        'y': np.arange(100),
        'x': np.arange(100),
        't': np.arange(10)
        }
    )

window_size = 10
# The actual rolling.
# Note, construct() returns the actual data in each window
# as it's passed over the images. This results in a DA which
# has a shape (y, x, t, yw, xw) - e.g., for each coord there
# is the x-y window (10x10). You have to pass the dict to
# construct so that it knows what to call the additional
# dimensions.
windows = (
    ds['image]
    .rolling(y=window_size, x=window_size, centre=True)
    .construct({'y': 'yw', 'x': 'xw')
    )
# Then we can just take the mean over the dims we want!
win_mean = windows.mean(dim=('yw', 'xw', 't'))
# Note, you could do this as a single chained command
# but it's clearer to see what's going on this way.