[Refactor]: Consider using flox to improve temporal averaging grouping logic

[tomvothecoder]

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:

1. One can only group by a single variable.
2. 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:

xcdat/xcdat/temporal.py

Lines 1266 to 1322 in c9bcbcd

	def _label_time_coords(self, time_coords: xr.DataArray) -> xr.DataArray:
	"""Labels time coordinates with a group for grouping.
	This methods labels time coordinates for grouping by first extracting
	specific xarray datetime components from time coordinates and storing
	them in a pandas DataFrame. After processing (if necessary) is performed
	on the DataFrame, it is converted to a numpy array of datetime
	objects. This numpy serves as the data source for the final
	DataArray of labeled time coordinates.
	Parameters
	----------
	time_coords : xr.DataArray
	The time coordinates.
	Returns
	-------
	xr.DataArray
	The DataArray of labeled time coordinates for grouping.
	Examples
	--------
	Original daily time coordinates:
	>>> <xarray.DataArray 'time' (time: 4)>
	>>> array(['2000-01-01T12:00:00.000000000',
	>>> '2000-01-31T21:00:00.000000000',
	>>> '2000-03-01T21:00:00.000000000',
	>>> '2000-04-01T03:00:00.000000000'],
	>>> dtype='datetime64[ns]')
	>>> Coordinates:
	>>> * time (time) datetime64[ns] 2000-01-01T12:00:00 ... 2000-04-01T03:00:00
	Daily time coordinates labeled by year and month:
	>>> <xarray.DataArray 'time' (time: 3)>
	>>> array(['2000-01-01T00:00:00.000000000',
	>>> '2000-03-01T00:00:00.000000000',
	>>> '2000-04-01T00:00:00.000000000'],
	>>> dtype='datetime64[ns]')
	>>> Coordinates:
	>>> * time (time) datetime64[ns] 2000-01-01T00:00:00 ... 2000-04-01T00:00:00
	"""
	df_dt_components: pd.DataFrame = self._get_df_dt_components(time_coords)
	dt_objects = self._convert_df_to_dt(df_dt_components)
	time_grouped = xr.DataArray(
	name="_".join(df_dt_components.columns),
	data=dt_objects,
	coords={self.dim: time_coords[self.dim]},
	dims=[self.dim],
	attrs=time_coords[self.dim].attrs,
	)
	time_grouped.encoding = time_coords[self.dim].encoding
	return time_grouped

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

1. Preprocess time coordinates (e.g., drop leap days, subset based on reference climatology)
2. 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)
3. Convert DataFrame to cftime objects to represent new time coordinates
4. Replace existing time coordinates in the DataArray with new time coordinates
5. 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 Multiple groupers v3 xarray-contrib/flox#76 is released in a new xarray version in Update GroupBy constructor for grouping by multiple variables, dask arrays pydata/xarray#6610, we should be able to do this.
• Also, Enable flox in GroupBy and resample pydata/xarray#5734 is now merged which improves .groupby() performance significantly.

• Support multi-dimensional grouped operations and group_over pydata/xarray#324 (comment)
• calculating climatologies efficiently pangeo-data/pangeo#271
• https://stackoverflow.com/questions/37008103/python-xarray-grouping-by-multiple-parameters
• https://stackoverflow.com/questions/54776283/how-to-call-the-xarrays-groupby-function-to-group-data-by-a-combination-of-year
• https://stackoverflow.com/questions/69784076/xarray-groupby-according-to-multi-indexs

[dcherian]

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

[tomvothecoder]

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