Update GroupBy constructor for grouping by multiple variables, dask arrays

[dcherian]

What is your issue?

flox supports grouping by multiple variables (would fix #324, #1056) and grouping by dask variables (would fix #2852).

To enable this in GroupBy we need to update the constructor's signature to

1. Accept multiple "by" variables.
2. Accept "expected group labels" for grouping by dask variables (like bins for groupby_bins which already supports grouping by dask variables). This lets us construct the output coordinate without evaluating the dask variable.
3. We may also want to simultaneously group by a categorical variable (season) and bin by a continuous variable (air temperature). So we also need a way to indicate whether the "expected group labels" are "bin edges" or categories.

---

The signature in flox is (may be errors!)

xarray_reduce(
    obj: Dataset | DataArray,
    *by: DataArray | str,
    func: str | Aggregation,
    expected_groups: Sequence | np.ndarray | None = None,
    isbin: bool | Sequence[bool] = False,
    ...
)

You would calculate that last example using flox as

xarray_reduce(
   ds,
    "season", "air_temperature", 
    expected_groups=[None, np.arange(21, 30, 1)],
    isbin=[False, True],
    ...
)

The use of expected_groups and isbin seems ugly to me (the names could also be better!)

---

I propose we update groupby's signature to

1. change group: DataArray | str to group: DataArray | str | Iterable[str] | Iterable[DataArray]
2. We could add a top-level xr.Bins object that wraps bin edges + any kwargs to be passed to pandas.cut. Note our current groupby_bins signature has a bunch of kwargs passed directly to pandas.cut.
3. Finally add groups: None | ArrayLike | xarray.Bins | Iterable[None | ArrayLike | xarray.Bins] to pass the "expected group labels".
   1. If None, then groups will be auto-detected from non-dask group arrays (if None for a dask group, then raise error).
   2. If xarray.Bins indicates binning by the appropriate variables
   3. If ArrayLike treat as categorical.
   4. groups is a little too similar to group so we should choose a better name.
   5. The ordering of ArrayLike would let us fix Ordered Groupby Keys #757 (pass the seasons in the order you want them in the output)

So then that example becomes

ds.groupby(
    ["season", "air_temperature"], # season is numpy, air_temperature is dask
    groups=[None, xr.Bins(np.arange(21, 30, 1), closed="right")],
)

Thoughts?

[malmans2]

I'm getting errors with multi-indexes and flox. Is this expected and related to this issue, or should I open a separate issue?

import numpy as np

import xarray as xr

ds = xr.Dataset(
    dict(a=(("z",), np.ones(10))),
    coords=dict(b=(("z"), np.arange(2).repeat(5)), c=(("z"), np.arange(5).repeat(2))),
).set_index(bc=["b", "c"])
grouped = ds.groupby("bc")

with xr.set_options(use_flox=False):
    grouped.sum()  # OK

with xr.set_options(use_flox=True):
    grouped.sum()  # Error

Traceback (most recent call last):
  File "/Users/mattia/MyGit/test.py", line 15, in <module>
    grouped.sum()
  File "/Users/mattia/MyGit/xarray/xarray/core/_reductions.py", line 2763, in sum
    return self._flox_reduce(
  File "/Users/mattia/MyGit/xarray/xarray/core/groupby.py", line 661, in _flox_reduce
    result = xarray_reduce(
  File "/Users/mattia/mambaforge/envs/sarsen_dev/lib/python3.10/site-packages/flox/xarray.py", line 373, in xarray_reduce
    actual[k] = v.expand_dims(missing_group_dims)
  File "/Users/mattia/MyGit/xarray/xarray/core/dataset.py", line 1427, in __setitem__
    self.update({key: value})
  File "/Users/mattia/MyGit/xarray/xarray/core/dataset.py", line 4432, in update
    merge_result = dataset_update_method(self, other)
  File "/Users/mattia/MyGit/xarray/xarray/core/merge.py", line 1070, in dataset_update_method
    return merge_core(
  File "/Users/mattia/MyGit/xarray/xarray/core/merge.py", line 722, in merge_core
    aligned = deep_align(
  File "/Users/mattia/MyGit/xarray/xarray/core/alignment.py", line 824, in deep_align
    aligned = align(
  File "/Users/mattia/MyGit/xarray/xarray/core/alignment.py", line 761, in align
    aligner.align()
  File "/Users/mattia/MyGit/xarray/xarray/core/alignment.py", line 550, in align
    self.assert_unindexed_dim_sizes_equal()
  File "/Users/mattia/MyGit/xarray/xarray/core/alignment.py", line 450, in assert_unindexed_dim_sizes_equal
    raise ValueError(
ValueError: cannot reindex or align along dimension 'bc' because of conflicting dimension sizes: {10, 6} (note: an index is found along that dimension with size=10)

[dcherian]

In xarray-contrib/flox#191 @keewis proposes a much nicer API for multiple variables:

data.groupby(
    xr.Grouper(by="x", bins=pd.IntervalIndex.from_breaks(coords["x_vertices"])),  # binning
    xr.Grouper(by=data.y, labels=["a", "b", "c"]),  # categorical, data.y is dask-backed
    xr.Grouper(by="time", freq="MS"),  # resample
)

Note pd.Grouper uses key instead of by so that's a possibility too.

[TomNicholas]

Using xr.Grouper has the advantage that you don't have to start guessing about whether or not the user wanted some complicated behaviour (especially if their input is slightly wrong somehow and you have to raise an informative error). Simple defaults would get left as is and complex use cases can be explicit and opt-in.

[shoyer]

I also like the idea of creating specific Grouper objects for different types of selection, e.g., UniqueGrouper (the default), BinGrouper, TimeResampleGrouper, etc.

[dcherian]

Here's a question.

In #7561, I implement Grouper objects that don't have any information of the variable we're grouping by. So the future API would be:

data.groupby({
	"x0": xr.BinGrouper(bins=pd.IntervalIndex.from_breaks(coords["x_vertices"])),  # binning
    "y": xr.UniqueGrouper(labels=["a", "b", "c"]),  # categorical, data.y is dask-backed
    "time": xr.TimeResampleGrouper(freq="MS")
	},
)

Does this look OK or do we want to support passing the DataArray or variable name as a by kwarg:

xr.BinGrouper(by="x0", bins=pd.IntervalIndex.from_breaks(coords["x_vertices"]))

This syntax would support passing DataArray in by so xr.UniqueGrouper(by=data.y) for example. Is that an important usecase to support? In #7561, I create new ResolvedGrouper objects that do contain by as a DataArray always, so it's really a question of exposing that to the user.

PS: Pandas has a key kwarg for a column name. So following that would mean

data.groupby([
	xr.BinGrouper("x0", bins=pd.IntervalIndex.from_breaks(coords["x_vertices"])),  # binning
    xr.UniqueGrouper("y", labels=["a", "b", "c"]),  # categorical, data.y is dask-backed
    xr.TimeResampleGrouper("time", freq="MS")
	],
)

[dcherian]

We voted to move forward with this API:

data.groupby({
	"x0": xr.BinGrouper(bins=pd.IntervalIndex.from_breaks(coords["x_vertices"])),  # binning
    "y": xr.UniqueGrouper(labels=["a", "b", "c"]),  # categorical, data.y is dask-backed
    "time": xr.TimeResampleGrouper(freq="MS")
	},
)

We won't break backwards-compatibility for da.groupby(other_data_array) but for any complicated use-cases with Grouper the user must add the by variable to the xarray object, and refer to it by name in the dictionary as above,