GroupBy(multiple groupers)

[dcherian]

Builds on #9389

• ClosesWIP: progress toward making groupby work with multiple arguments #924,
• Closesgroupby_bins along two dims simultaneously #1056,
• ClosesEnable multi-coord grouping from xarray #9332,
• xref Support group_over #324
• Tests added
• User visible changes (including notable bug fixes) are documented in whats-new.rst
• New functions/methods are listed in api.rst
• add narrative docs
• Update docstrings

---

This PR takes an approach similar to that in #924.

The GroupBy(multiple array) problem can be reduced to the GroupBy(single nD variable) problem:

1. First convert each by array to integer codes ("factorize"). I do this first because factorization is expensive, and so it's cheaper to do it before any possible broadcasting.
2. Broadcast DataArrays of integer codes against each other.
3. Then np.ravel_multi_index them to construct a single nD array of integer codes.
4. At this point, we treat it as a GroupBy(single nD array) problem i.e. reshape down to 1D problem. This isn't great for dask/cubed but it is simple and matches what we do today. That means we can apply all the operations we do with GroupBy today.
   • Importantly UDFs receive "stacked arrays" as they do today when grouping by a nD array. This isn't the most elegant but it preserves dtype. For example, you can't index out a 2D array with group==1 without inserting a missing value if group is:

   [[1, 1, 1],
   [2, 1, 1],
   [2, 2, 1]]
   

5. Once operations are complete, we assign a new MultiIndex and then unstack back to restore the needed dimensions.
   • I'm not sure this is the best at the moment - what if one of the arrays we group by is a multiIndexed array?

TODO:

• lots more tests!

[dcherian]

@max-sixty I could use some help with adding tests,docs,docstrings if you have time :)

[dcherian]

This is the main new logic.

[max-sixty]

I had a look at where I add something — it already seems extremely good? gb - gb.mean() would be cool but this already seems like a big step forwards.

I missed previous discussions / can decide this later — but do we want to have da.groupby(['foo', 'bar']) as sugar for da.groupby(foo=UniqueGrouper(), bar=UniqueGrouper())?

[dcherian]

Yes it is surprisingly complete.

A section in the narrative docs would be nice to add.

[dcherian]

missed previous discussions / can decide this later — but do we want to have da.groupby(['foo', 'bar']) as sugar for da.groupby(foo=UniqueGrouper(), bar=UniqueGrouper())?

I agree. It's a bit too much typing at the moment

[dcherian]

We should test grouping by a multi-index array and some other array but otherwise this should be good to go.

EDIT: Now raises NotImplementedError for this case.

[max-sixty]

This is huge! Thank you very much @dcherian !

[dcherian]

Merging so it's ready to try out!

[max-sixty]

Huge success! Thank you @dcherian !

[shoyer]

Wow, amazing work @dcherian! I am so excited to see this finally land.