Add GroupBy.get_group. copying Pandas API

Author: dcherian

doc/api.rst

@@ -559,9 +559,11 @@ GroupBy objects
    core.groupby.DataArrayGroupBy
    core.groupby.DataArrayGroupBy.map
    core.groupby.DataArrayGroupBy.reduce
+   core.groupby.DataArrayGroupBy.get_group
    core.groupby.DatasetGroupBy
    core.groupby.DatasetGroupBy.map
    core.groupby.DatasetGroupBy.reduce
+   core.groupby.DatasetGroupBy.get_group
 
 Rolling objects
 ===============

doc/whats-new.rst

@@ -34,6 +34,10 @@ New Features
 - :py:meth:`Dataset.quantile`, :py:meth:`DataArray.quantile` and ``GroupBy.quantile``
   now work with dask Variables.
   By `Deepak Cherian <https://github.com/dcherian>`_.
+- Implement ``get_group`` 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>`_.
 - Added the ``count`` reduction method to both :py:class:`~core.rolling.DatasetCoarsen`
   and :py:class:`~core.rolling.DataArrayCoarsen` objects. (:pull:`3500`)
   By `Deepak Cherian <https://github.com/dcherian>`_

xarray/core/groupby.py

@@ -426,11 +426,24 @@ 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 get_group(self, name):
+        """
+        Get DataArray or Dataset corresponding to a particular group label.
+        """
+        if name not in self.groups:
+            raise KeyError(f"{name} is not a valid group label.")
+
+        indices = self.groups[name]
+        return self._obj.isel(**{self._group_dim: indices})
+
     def __len__(self):
         return self._unique_coord.size
 

xarray/tests/test_groupby.py

@@ -499,4 +499,17 @@ def test_groupby_bins_timeseries():
     assert_identical(actual, expected)
 
 
+def test_groupby_get_group(dataset):
+
+    assert_identical(dataset.sel(x="a"), dataset.groupby("x").get_group("a"))
+    assert_identical(dataset.sel(z=1), dataset.groupby("z").get_group(1))
+
+    assert_identical(dataset.foo.sel(x="a"), dataset.foo.groupby("x").get_group("a"))
+    assert_identical(dataset.foo.sel(z=1), dataset.foo.groupby("z").get_group(1))
+
+    actual = dataset.groupby("boo").get_group("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