Using flox to resample

[tlsw231]

I want to see if using flox improves the performance of some resample operations I am trying on a large dataset. The data is 3-hourly and I want to resample to daily frequency.

1. I came across this link which demonstrates usage of flox.xarray.resample_reduce but this method is missing from the docs. Any particular reason?

2. Here is a dataset with the same size as my original one (1.5 TB). There is some reduction in the number of tasks with flox. Is it possible to do even better?

import xarray as xr

import flox
import flox.xarray
from flox.xarray import xarray_reduce

import dask
from dask.distributed import Client
from dask import delayed # to allow parallel computation
from distributed import Client
from dask_jobqueue import SLURMCluster

# Request cores
ncores=8; nprocs=ncores; totalcores= ncores * 1
cluster = SLURMCluster(cores=ncores, processes=nprocs, memory="264GB",queue="medium",
                       walltime="00:30:00")
cluster.scale(totalcores)
client = Client(cluster)

# Create dataarray of shape (23368, 28, 489, 655)
times = xr.cftime_range(start = '2011-01-01 03:00:00', end = '2018-12-31', freq='3H')
(nt,nz,ny,nx) =(len(times),28,489,655)
da = xr.DataArray(data=dask.array.random.random([nt,nz,ny,nx]),
                  dims=['time','z','y','x'],coords=dict(time=times)).rename('dummy')

Here is the data:

Using the standard xarray method:

da.resample(time="1D").mean()

Using flox:

flox.xarray.resample_reduce(da.resample(time="1D"), func="mean")

[dcherian]

Thanks. This seems to be the difference between method="cohorts" and method="blockwise", the latter is hardcoded in resample_reduce.

With blockwise we rechunk so that resampling becomes a blockwise embarassingly parallel problem. It used to be the default for da.resample but I forget why I changed it now.

So

da.resample(time="1D").mean(method="blockwise")

should recover the resample_reduce behaviour, which is why resample_reduce was removed. (I'll add a DeprecationWarning)

[dcherian]

Also this problem is really sensitive to chunk size so dask.array.random.random([nt,nz,ny,nx]) (basically auto chunking) is not a good idea. What is the chunking along time for the original problem? Since there are 8 3H time steps in 1D a multiple of 8 will be best and make resampling a blockwise problem.