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 563be73ddea..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 diff --git a/xarray/core/groupby.py b/xarray/core/groupby.py index 824f2767153..8a0972519e1 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..5ef7677c499 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_getitem(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