Add DatasetGroupBy.quantile

[keewis]

This adds DatasetGroupBy.quantile by moving DataArrayGroupBy.quantile to GroupBy as proposed in #3018 (comment).

The tests are a modified copy of the ones from #2828. What confuses me is that expected_yy in test_ds_groupby_quantile needs the transpose whereas the equivalent in test_da_groupby_quantile doesn't. Does anyone have an idea about why that is?

• Closes DatasetGroupBy does not implement quantile #3525
• Tests added
• Passes black . && mypy . && flake8
• Fully documented, including whats-new.rst for all changes and api.rst for new API

[keewis]

is Bug Fixes the correct section for this?

[dcherian]

Let's do enhancements!

[keewis]

that would open a new category. Should I use New Features?

[dcherian]

👍

[dcherian]

Thanks @keewis !

I left some minor comments. The transpose is surprising. It would be good to track that down to either .map or Groupby._combine

[keewis]

the difference is here, where in DataArrayGroupBy the original dimension order is restored. DatasetGroupBy does not do that, so the test requires the transpose

[dcherian]

LGTM @keewis.

Re: docstring. If you jave time, can you add a couple of examples (just pick two from the tests) to illustrate the scalar vs vector behaviour of q?

[keewis]

sure. But as Dataset.quantile and DataArray.quantile don't have examples either, I think it might be worth to do that in a new PR?

Also, what will we do re transpose?

[dcherian]

👍

The transpose is fine, I think. It looks intentional and must be consistent with other Dataset/DataArray operations?

[keewis]

well, it was surprising to me. It also might be problematic that it matters if we first pull out a DataArray and then compute or the other way around:

>>> ds = xr.Dataset(
...     data_vars={
...         "a": (
...             ("x", "y"),
...             [[1, 11, 26], [2, 12, 22], [3, 13, 23],^I[4, 16,^I24], [5, 15, 25]],
...         )
...     },
...     coords={"x": [1, 1, 1, 2, 2], "y": [0, 0, 1]},
... )
>>> a_ds = ds.groupby("y").quantile(0).a
>>> a_da = ds.a.groupby("y").quantile(0)
>>> a_ds.identical(a_da)
False
>>> a_ds.transpose().identical(a_da)
True

[dcherian]

Hmm.. ok calling mean in this example passes the identical check so maybe there's something funny in quantile:

ds = xr.Dataset(
    data_vars={
        "a": (
            ("x", "y"),
            [[1, 11, 26], [2, 12, 22], [3, 13, 23],
             [4, 16,24], [5, 15, 25]],
        )
    },
    coords={"x": [1, 1, 1, 2, 2], "y": [0, 0, 1]},
).expand_dims({"z": [0,1,2]})
a_ds = ds.groupby("y").mean("x").a
a_da = ds.a.groupby("y").mean("x")
a_ds.identical(a_da)

[keewis]

but that's only because mean removes x, doesn't it?

[keewis]

I can also reproduce this with std

Edit: I'm fairly certain this is due to what I mentioned in #3527 (comment)

[dcherian]

I can also reproduce this with std

Which behaviour can you reproduce?

Replacing mean with std in my last example passes the identical check. Actually this works too

a_ds = ds.groupby("y").apply(lambda x: x).a
a_da = ds.a.groupby("y").apply(lambda x: x)
a_ds.identical(a_da)

[max-sixty]

Dimension reordering has come up a few times before (#1739, and others that I can't immediately find)

I think we should probably have a general solution & standard behavior here. Others should weigh in (CC @shoyer), but I would vote to not worry too much about it here, and prioritize simplicity and consistency with past behavior until we decide on a more general solution.

[shoyer]

the difference is here, where in DataArrayGroupBy the original dimension order is restored. DatasetGroupBy does not do that, so the test requires the transpose

It was not entirely obvious to me how to generalize "restoring dimension order" to Dataset. On a Dataset, it is not entirely clear which variable dimension order should be copied from. Maybe each variable in the result should have dimension order copied from the variable with the same name in the original dataset? But then what about new variables?

I agree @max-sixty, let's not worry about this too much for now.

[dcherian]

In it goes. Great work as usual, @keewis. Thanks!

[keewis]

Which behaviour can you reproduce?

ds = xr.Dataset(
    data_vars={
        "a": (
            ("x", "y"),
            [[1, 11, 26], [2, 12, 22], [3, 13, 23],
             [4, 16,24], [5, 15, 25]],
        )
    },
    coords={"x": [1, 1, 1, 2, 2], "y": [0, 0, 1]},
).expand_dims({"z": [0, 1, 1, 2, 2]})
a_ds = ds.groupby("y").std().a
a_da = ds.a.groupby("y").std()
a_ds.identical(a_da)

which seems to be the same as

a_ds = ds.groupby("y").mean("y").a
a_da = ds.a.groupby("y").mean("y")
a_ds.identical(a_da)

My impression is that this puts the dimension grouped over to the front:

• before: ds.a.dims == ("z", "x", "y")
• afterwards: a_ds.dims == ("y", "z", "x")

But I agree this discussion does not really belong here. So, thanks all!