From ed8d2263110c6b7a393eb6ae423b004925856118 Mon Sep 17 00:00:00 2001 From: dcherian Date: Sat, 6 Feb 2021 09:07:23 -0700 Subject: [PATCH 1/3] Implement GroupBy.__getitem__ --- doc/groupby.rst | 6 ++++++ doc/whats-new.rst | 4 ++++ xarray/core/groupby.py | 9 +++++++++ xarray/tests/test_groupby.py | 13 +++++++++++++ 4 files changed, 32 insertions(+) diff --git a/doc/groupby.rst b/doc/groupby.rst index d0c0b1849f9..804799aebd1 100644 --- a/doc/groupby.rst +++ b/doc/groupby.rst @@ -63,6 +63,12 @@ You can also iterate over groups in ``(label, group)`` pairs: list(ds.groupby("letters")) +You can index out a particular group: + +.. ipython:: python + + ds.groupby("letters")["b"] + Just like in pandas, creating a GroupBy object is cheap: it does not actually split the data until you access particular values. diff --git a/doc/whats-new.rst b/doc/whats-new.rst index 471e91a8512..158ddf3350a 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -62,6 +62,10 @@ New Features - :py:meth:`DataArray.swap_dims` & :py:meth:`Dataset.swap_dims` now accept dims in the form of kwargs as well as a dict, like most similar methods. By `Maximilian Roos `_. +- Implement ``__getitem__`` for both :py:class:`~core.groupby.DatasetGroupBy` and + :py:class:`~core.groupby.DataArrayGroupBy`, inspired by pandas' + :py:meth:`~pandas.core.groupby.GroupBy.get_group`. + By `Deepak Cherian `_. Bug fixes ~~~~~~~~~ diff --git a/xarray/core/groupby.py b/xarray/core/groupby.py index e1e5a0fabe8..bfcbd4b62ee 100644 --- a/xarray/core/groupby.py +++ b/xarray/core/groupby.py @@ -415,11 +415,20 @@ def dims(self): @property def groups(self): + """ + Mapping from group labels to indices. The indices can be used to index the underlying object. + """ # provided to mimic pandas.groupby if self._groups is None: self._groups = dict(zip(self._unique_coord.values, self._group_indices)) return self._groups + def __getitem__(self, key): + """ + Get DataArray or Dataset corresponding to a particular group label. + """ + return self._obj.isel({self._group_dim: self.groups[key]}) + def __len__(self): return self._unique_coord.size diff --git a/xarray/tests/test_groupby.py b/xarray/tests/test_groupby.py index 85f729a9f7a..1366dd417f7 100644 --- a/xarray/tests/test_groupby.py +++ b/xarray/tests/test_groupby.py @@ -549,4 +549,17 @@ def test_groupby_none_group_name(): assert "group" in mean.dims +def test_groupby_sel(dataset): + + assert_identical(dataset.sel(x="a"), dataset.groupby("x")["a"]) + assert_identical(dataset.sel(z=1), dataset.groupby("z")[1]) + + assert_identical(dataset.foo.sel(x="a"), dataset.foo.groupby("x")["a"]) + assert_identical(dataset.foo.sel(z=1), dataset.foo.groupby("z")[1]) + + actual = dataset.groupby("boo")["f"].unstack().transpose("x", "y", "z") + expected = dataset.sel(y=[1], z=[1, 2]).transpose("x", "y", "z") + assert_identical(expected, actual) + + # TODO: move other groupby tests from test_dataset and test_dataarray over here From fe036ae443ecc202a04877b67526133a48963b43 Mon Sep 17 00:00:00 2001 From: dcherian Date: Sun, 7 Feb 2021 15:15:27 -0700 Subject: [PATCH 2/3] fix test name --- xarray/tests/test_groupby.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/tests/test_groupby.py b/xarray/tests/test_groupby.py index 1366dd417f7..5ef7677c499 100644 --- a/xarray/tests/test_groupby.py +++ b/xarray/tests/test_groupby.py @@ -549,7 +549,7 @@ def test_groupby_none_group_name(): assert "group" in mean.dims -def test_groupby_sel(dataset): +def test_groupby_getitem(dataset): assert_identical(dataset.sel(x="a"), dataset.groupby("x")["a"]) assert_identical(dataset.sel(z=1), dataset.groupby("z")[1]) From 8994c5b4ba9d3ef8f1f0bb33c3ce9147dd1e11df Mon Sep 17 00:00:00 2001 From: Mathias Hauser Date: Sat, 13 Mar 2021 20:23:04 +0100 Subject: [PATCH 3/3] move whats new entry --- doc/whats-new.rst | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index f961de61af5..c7fba22e4ad 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -51,6 +51,10 @@ New Features grant from the `Chan Zuckerberg Initiative `_ and developed by `B-Open `_. By `Aureliana Barghini `_ and `Alessandro Amici `_. +- Implement ``__getitem__`` for both :py:class:`~core.groupby.DatasetGroupBy` and + :py:class:`~core.groupby.DataArrayGroupBy`, inspired by pandas' + :py:meth:`~pandas.core.groupby.GroupBy.get_group`. + By `Deepak Cherian `_. Breaking changes @@ -197,10 +201,6 @@ New Features - :py:meth:`DataArray.swap_dims` & :py:meth:`Dataset.swap_dims` now accept dims in the form of kwargs as well as a dict, like most similar methods. By `Maximilian Roos `_. -- Implement ``__getitem__`` for both :py:class:`~core.groupby.DatasetGroupBy` and - :py:class:`~core.groupby.DataArrayGroupBy`, inspired by pandas' - :py:meth:`~pandas.core.groupby.GroupBy.get_group`. - By `Deepak Cherian `_. Bug fixes ~~~~~~~~~