Grouper, Resampler as public api (#8840)

* Grouper, Resampler as public API * Add test * Add docs * Fix test * fix types. * bugfix * Better binning API * docs.fixes * Apply suggestions from code review Co-authored-by: Illviljan <[email protected]> * Fix typing * clean up reprs * Allow passing dicts * Apply suggestions from code review Co-authored-by: Justus Magin <[email protected]> * Update xarray/core/common.py Co-authored-by: Justus Magin <[email protected]> * Review comments * Fix docstring * Try to fix typing * Nicer error * Try fixing types * fix * Apply suggestions from code review Co-authored-by: Michael Niklas <[email protected]> * Review comments * Add whats-new note * Fix * Add more types * Fix link --------- Co-authored-by: Illviljan <[email protected]> Co-authored-by: Justus Magin <[email protected]> Co-authored-by: Michael Niklas <[email protected]>
pydata · Jul 18, 2024 · b55c783 · b55c783
1 parent c2aebd8
commit b55c783
Show file tree

Hide file tree

Showing 13 changed files with 477 additions and 160 deletions.
diff --git a/doc/api-hidden.rst b/doc/api-hidden.rst
@@ -693,3 +693,7 @@
 
    coding.times.CFTimedeltaCoder
    coding.times.CFDatetimeCoder
+
+   core.groupers.Grouper
+   core.groupers.Resampler
+   core.groupers.EncodedGroups
diff --git a/doc/api.rst b/doc/api.rst
@@ -803,6 +803,18 @@ DataArray
    DataArrayGroupBy.dims
    DataArrayGroupBy.groups
 
+Grouper Objects
+---------------
+
+.. currentmodule:: xarray.core
+
+.. autosummary::
+   :toctree: generated/
+
+   groupers.BinGrouper
+   groupers.UniqueGrouper
+   groupers.TimeResampler
+
 
 Rolling objects
 ===============
@@ -1028,17 +1040,20 @@ DataArray
 Accessors
 =========
 
-.. currentmodule:: xarray
+.. currentmodule:: xarray.core
 
 .. autosummary::
    :toctree: generated/
 
-   core.accessor_dt.DatetimeAccessor
-   core.accessor_dt.TimedeltaAccessor
-   core.accessor_str.StringAccessor
+   accessor_dt.DatetimeAccessor
+   accessor_dt.TimedeltaAccessor
+   accessor_str.StringAccessor
+
 
 Custom Indexes
 ==============
+.. currentmodule:: xarray
+
 .. autosummary::
    :toctree: generated/
 

diff --git a/doc/conf.py b/doc/conf.py
@@ -158,6 +158,8 @@
     "Variable": "~xarray.Variable",
     "DatasetGroupBy": "~xarray.core.groupby.DatasetGroupBy",
     "DataArrayGroupBy": "~xarray.core.groupby.DataArrayGroupBy",
+    "Grouper": "~xarray.core.groupers.Grouper",
+    "Resampler": "~xarray.core.groupers.Resampler",
     # objects without namespace: numpy
     "ndarray": "~numpy.ndarray",
     "MaskedArray": "~numpy.ma.MaskedArray",
@@ -169,6 +171,7 @@
     "CategoricalIndex": "~pandas.CategoricalIndex",
     "TimedeltaIndex": "~pandas.TimedeltaIndex",
     "DatetimeIndex": "~pandas.DatetimeIndex",
+    "IntervalIndex": "~pandas.IntervalIndex",
     "Series": "~pandas.Series",
     "DataFrame": "~pandas.DataFrame",
     "Categorical": "~pandas.Categorical",

diff --git a/doc/user-guide/groupby.rst b/doc/user-guide/groupby.rst
@@ -1,3 +1,5 @@
+.. currentmodule:: xarray
+
 .. _groupby:
 
 GroupBy: Group and Bin Data
@@ -15,19 +17,20 @@ __ https://www.jstatsoft.org/v40/i01/paper
 - Apply some function to each group.
 - Combine your groups back into a single data object.
 
-Group by operations work on both :py:class:`~xarray.Dataset` and
-:py:class:`~xarray.DataArray` objects. Most of the examples focus on grouping by
+Group by operations work on both :py:class:`Dataset` and
+:py:class:`DataArray` objects. Most of the examples focus on grouping by
 a single one-dimensional variable, although support for grouping
 over a multi-dimensional variable has recently been implemented. Note that for
 one-dimensional data, it is usually faster to rely on pandas' implementation of
 the same pipeline.
 
 .. tip::
 
-   To substantially improve the performance of GroupBy operations, particularly
-   with dask `install the flox package <https://flox.readthedocs.io>`_. flox
+   `Install the flox package <https://flox.readthedocs.io>`_ to substantially improve the performance
+   of GroupBy operations, particularly with dask. flox
    `extends Xarray's in-built GroupBy capabilities <https://flox.readthedocs.io/en/latest/xarray.html>`_
-   by allowing grouping by multiple variables, and lazy grouping by dask arrays. If installed, Xarray will automatically use flox by default.
+   by allowing grouping by multiple variables, and lazy grouping by dask arrays.
+   If installed, Xarray will automatically use flox by default.
 
 Split
 ~~~~~
@@ -87,7 +90,7 @@ Binning
 Sometimes you don't want to use all the unique values to determine the groups
 but instead want to "bin" the data into coarser groups. You could always create
 a customized coordinate, but xarray facilitates this via the
-:py:meth:`~xarray.Dataset.groupby_bins` method.
+:py:meth:`Dataset.groupby_bins` method.
 
 .. ipython:: python
 
@@ -110,7 +113,7 @@ Apply
 ~~~~~
 
 To apply a function to each group, you can use the flexible
-:py:meth:`~xarray.core.groupby.DatasetGroupBy.map` method. The resulting objects are automatically
+:py:meth:`core.groupby.DatasetGroupBy.map` method. The resulting objects are automatically
 concatenated back together along the group axis:
 
 .. ipython:: python
@@ -121,8 +124,8 @@ concatenated back together along the group axis:
 
     arr.groupby("letters").map(standardize)
 
-GroupBy objects also have a :py:meth:`~xarray.core.groupby.DatasetGroupBy.reduce` method and
-methods like :py:meth:`~xarray.core.groupby.DatasetGroupBy.mean` as shortcuts for applying an
+GroupBy objects also have a :py:meth:`core.groupby.DatasetGroupBy.reduce` method and
+methods like :py:meth:`core.groupby.DatasetGroupBy.mean` as shortcuts for applying an
 aggregation function:
 
 .. ipython:: python
@@ -183,7 +186,7 @@ Iterating and Squeezing
 Previously, Xarray defaulted to squeezing out dimensions of size one when iterating over
 a GroupBy object. This behaviour is being removed.
 You can always squeeze explicitly later with the Dataset or DataArray
-:py:meth:`~xarray.DataArray.squeeze` methods.
+:py:meth:`DataArray.squeeze` methods.
 
 .. ipython:: python
 
@@ -217,7 +220,7 @@ __ https://cfconventions.org/cf-conventions/v1.6.0/cf-conventions.html#_two_dime
     da.groupby("lon").map(lambda x: x - x.mean(), shortcut=False)
 
 Because multidimensional groups have the ability to generate a very large
-number of bins, coarse-binning via :py:meth:`~xarray.Dataset.groupby_bins`
+number of bins, coarse-binning via :py:meth:`Dataset.groupby_bins`
 may be desirable:
 
 .. ipython:: python
@@ -232,3 +235,71 @@ applying your function, and then unstacking the result:
 
     stacked = da.stack(gridcell=["ny", "nx"])
     stacked.groupby("gridcell").sum(...).unstack("gridcell")
+
+.. _groupby.groupers:
+
+Grouper Objects
+~~~~~~~~~~~~~~~
+
+Both ``groupby_bins`` and ``resample`` are specializations of the core ``groupby`` operation for binning,
+and time resampling. Many problems demand more complex GroupBy application: for example, grouping by multiple
+variables with a combination of categorical grouping, binning, and resampling; or more specializations like
+spatial resampling; or more complex time grouping like special handling of seasons, or the ability to specify
+custom seasons. To handle these use-cases and more, Xarray is evolving to providing an
+extension point using ``Grouper`` objects.
+
+.. tip::
+
+   See the `grouper design`_ doc for more detail on the motivation and design ideas behind
+   Grouper objects.
+
+.. _grouper design: https://github.com/pydata/xarray/blob/main/design_notes/grouper_objects.md
+
+For now Xarray provides three specialized Grouper objects:
+
+1. :py:class:`groupers.UniqueGrouper` for categorical grouping
+2. :py:class:`groupers.BinGrouper` for binned grouping
+3. :py:class:`groupers.TimeResampler` for resampling along a datetime coordinate
+
+These provide functionality identical to the existing ``groupby``, ``groupby_bins``, and ``resample`` methods.
+That is,
+
+.. code-block:: python
+
+    ds.groupby("x")
+
+is identical to
+
+.. code-block:: python
+
+    from xarray.groupers import UniqueGrouper
+
+    ds.groupby(x=UniqueGrouper())
+
+; and
+
+.. code-block:: python
+
+    ds.groupby_bins("x", bins=bins)
+
+is identical to
+
+.. code-block:: python
+
+    from xarray.groupers import BinGrouper
+
+    ds.groupby(x=BinGrouper(bins))
+
+and
+
+.. code-block:: python
+
+    ds.resample(time="ME")
+
+is identical to
+
+.. code-block:: python
+
+    from xarray.groupers import TimeResampler
+
+    ds.resample(time=TimeResampler("ME"))
diff --git a/doc/whats-new.rst b/doc/whats-new.rst
@@ -22,6 +22,12 @@ v2024.06.1 (unreleased)
 
 New Features
 ~~~~~~~~~~~~
+- Introduce new :py:class:`groupers.UniqueGrouper`, :py:class:`groupers.BinGrouper`, and
+  :py:class:`groupers.TimeResampler` objects as a step towards supporting grouping by
+  multiple variables. See the `docs <groupby.groupers_>` and the
+  `grouper design doc <https://github.com/pydata/xarray/blob/main/design_notes/grouper_objects.md>`_ for more.
+  (:issue:`6610`, :pull:`8840`).
+  By `Deepak Cherian <https://github.com/dcherian>`_.
 - Allow per-variable specification of ``mask_and_scale``, ``decode_times``, ``decode_timedelta``
   ``use_cftime`` and ``concat_characters`` params in :py:func:`~xarray.open_dataset`  (:pull:`9218`).
   By `Mathijs Verhaegh <https://github.com/Ostheer>`_.

diff --git a/xarray/__init__.py b/xarray/__init__.py
@@ -14,6 +14,7 @@
 from xarray.coding.cftimeindex import CFTimeIndex
 from xarray.coding.frequencies import infer_freq
 from xarray.conventions import SerializationWarning, decode_cf
+from xarray.core import groupers
 from xarray.core.alignment import align, broadcast
 from xarray.core.combine import combine_by_coords, combine_nested
 from xarray.core.common import ALL_DIMS, full_like, ones_like, zeros_like
@@ -55,6 +56,7 @@
 # `mypy --strict` running in projects that import xarray.
 __all__ = (
     # Sub-packages
+    "groupers",
     "testing",
     "tutorial",
     # Top-level functions

diff --git a/xarray/core/common.py b/xarray/core/common.py
@@ -38,6 +38,7 @@
 
     from xarray.core.dataarray import DataArray
     from xarray.core.dataset import Dataset
+    from xarray.core.groupers import Resampler
     from xarray.core.indexes import Index
     from xarray.core.resample import Resample
     from xarray.core.rolling_exp import RollingExp
@@ -876,7 +877,7 @@ def rolling_exp(
     def _resample(
         self,
         resample_cls: type[T_Resample],
-        indexer: Mapping[Any, str] | None,
+        indexer: Mapping[Hashable, str | Resampler] | None,
         skipna: bool | None,
         closed: SideOptions | None,
         label: SideOptions | None,
@@ -885,7 +886,7 @@ def _resample(
         origin: str | DatetimeLike,
         loffset: datetime.timedelta | str | None,
         restore_coord_dims: bool | None,
-        **indexer_kwargs: str,
+        **indexer_kwargs: str | Resampler,
     ) -> T_Resample:
         """Returns a Resample object for performing resampling operations.
 
@@ -1068,7 +1069,7 @@ def _resample(
 
         from xarray.core.dataarray import DataArray
         from xarray.core.groupby import ResolvedGrouper
-        from xarray.core.groupers import TimeResampler
+        from xarray.core.groupers import Resampler, TimeResampler
         from xarray.core.resample import RESAMPLE_DIM
 
         # note: the second argument (now 'skipna') use to be 'dim'
@@ -1098,15 +1099,21 @@ def _resample(
             name=RESAMPLE_DIM,
         )
 
-        grouper = TimeResampler(
-            freq=freq,
-            closed=closed,
-            label=label,
-            origin=origin,
-            offset=offset,
-            loffset=loffset,
-            base=base,
-        )
+        grouper: Resampler
+        if isinstance(freq, str):
+            grouper = TimeResampler(
+                freq=freq,
+                closed=closed,
+                label=label,
+                origin=origin,
+                offset=offset,
+                loffset=loffset,
+                base=base,
+            )
+        elif isinstance(freq, Resampler):
+            grouper = freq
+        else:
+            raise ValueError("freq must be a str or a Resampler object")
 
         rgrouper = ResolvedGrouper(grouper, group, self)