Skip to content

Commit

Permalink
Implement GroupBy.__getitem__
Browse files Browse the repository at this point in the history
  • Loading branch information
dcherian committed Feb 6, 2021
1 parent 5735e16 commit ed8d226
Show file tree
Hide file tree
Showing 4 changed files with 32 additions and 0 deletions.
6 changes: 6 additions & 0 deletions doc/groupby.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down
4 changes: 4 additions & 0 deletions doc/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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 <https://github.com/max-sixty>`_.
- 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 <https://github.com/dcherian>`_.

Bug fixes
~~~~~~~~~
Expand Down
9 changes: 9 additions & 0 deletions xarray/core/groupby.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
13 changes: 13 additions & 0 deletions xarray/tests/test_groupby.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

0 comments on commit ed8d226

Please sign in to comment.