[Refactor]: Consider using `flox` and `xr.resample()` to improve temporal averaging grouping logic

tomvothecoder commented 2 years ago

Is your feature request related to a problem?

Currently, grouping by multiple coordinates (e.g., time.year and time.season) requires creating a new set of coordinates before grouping due to the xarray limitations described below.

Xarray's GroupBy operations are currently limited:

One can only group by a single variable.

When grouping by a dask array, that array will be computed to discover the unique group labels, and their locations

-- source: https://flox.readthedocs.io/en/latest/xarray.html

Related code in xcdat for temporal grouping: https://github.com/xCDAT/xcdat/blob/c9bcbcdb66af916958a79a33177bc43d478e4036/xcdat/temporal.py#L1266-L1322

Current temporal averaging logic (workaround for multi-variable grouping):

Preprocess time coordinates (e.g., drop leap days, subset based on reference climatology)
Transform time coordinates from an xarray.DataArray to a pandas.DataFrame, a. Keep only the DataFrame columns needed for grouping (e.g., "year" and "season" for seasonal group averages), essentially "labeling" coordinates with their groups b. Process the DataFrame including:
- Mapping of months to custom seasons for custom seasonal grouping
- Correction of "DJF" seasons by shifting Decembers over to the next year
- Mapping of seasons to their mid months to create cftime coordinates (season strings aren't supported in cftime/datetime objects)
Convert DataFrame to cftime objects to represent new time coordinates
Replace existing time coordinates in the DataArray with new time coordinates
Group DataArray with new time coordinates for the mean

Describe the solution you'd like

It is would be simpler, cleaner, and probably more performant to call something like .groupby(["time.year", "time.season"]) instead (waiting on xarray to support this with flox). This solution will reduce a lot of the internal complexities involved with the temporal averaging API.

We might able to achieve this using flox directly:

These limitations can be avoided by using {py:func}flox.xarray.xarray_reduce which allows grouping by multiple variables, lazy grouping by dask variables, as well as an arbitrary combination of categorical grouping and binning. For example,
flox.xarray.xarray_reduce(
    ds,
    ds.time.dt.month,
    ds.lon,
    func="mean",
    expected_groups=[None, [0, 10, 20]],
    isbin=[False, True],
    method="map-reduce",
)
-- source: https://flox.readthedocs.io/en/latest/xarray.html

Additionally, would need to figure out a way to easily perform the processing steps for time coordinates directly in xarray objects described in 2b if we move away from using pandas.DataFrame.

Describe alternatives you've considered

Multi-variable grouping was originally done using pd.MultiIndex but we shifted away from this approach because this object cannot be written out to netcdf4. Also pd.MultiIndex is not the standard object type for representing time coordinates in xarray. The standard object types are np.datetime64 and cftime.

Additional context

Future solution through xarray + flox:

Once https://github.com/dcherian/flox/pull/76 is released in a new xarray version in https://github.com/pydata/xarray/issues/6610, we should be able to do this.
Also, https://github.com/pydata/xarray/pull/5734 is now merged which improves .groupby() performance significantly.

dcherian commented 1 year ago

I saw the ping at https://github.com/pydata/xarray/issues/6610. Let me know if you run in to issues or have questions

tomvothecoder commented 1 year ago

Thanks @dcherian! I'm looking forward to trying out flox.

tomvothecoder commented 3 weeks ago

xarray >= 2024.09.0 now supports grouping by multiple variables: https://xarray.dev/blog/multiple-groupers

xCDAT / xcdat