Implement groupby.weighted

Author: dcherian

xarray/core/groupby.py

@@ -5,6 +5,7 @@
 import pandas as pd
 
 from . import dtypes, duck_array_ops, nputils, ops
+from .alignment import align
 from .arithmetic import DataArrayGroupbyArithmetic, DatasetGroupbyArithmetic
 from .concat import concat
 from .formatting import format_array_flat
@@ -258,6 +259,7 @@ class GroupBy:
         "_stacked_dim",
         "_unique_coord",
         "_dims",
+        "_weighted",
     )
 
     def __init__(
@@ -395,6 +397,7 @@ def __init__(
         self._inserted_dims = inserted_dims
         self._full_index = full_index
         self._restore_coord_dims = restore_coord_dims
+        self._weighted = False
 
         # cached attributes
         self._groups = None
@@ -431,12 +434,32 @@ def __len__(self):
     def __iter__(self):
         return zip(self._unique_coord.values, self._iter_grouped())
 
+    def weighted(self, weights):
+        """Set weights for GroupBy reductions."""
+        self._weighted = True
+
+        # exact alignment to ensure we don't add NaNs to the weights
+        align(self._obj, weights, join="exact")
+
+        # Ensures that weights will get sliced appropriately
+        self._obj.coords["__weights__"] = weights
+
+        # save the name for repr
+        self._obj.coords["__weights__"].attrs["__name__"] = weights.name
+
+        return self
+
     def __repr__(self):
-        return "{}, grouped over {!r}\n{!r} groups with labels {}.".format(
-            self.__class__.__name__,
-            self._unique_coord.name,
-            self._unique_coord.size,
-            ", ".join(format_array_flat(self._unique_coord, 30).split()),
+        if self._weighted:
+            weights_name = self._obj.__weights__.attrs["__name__"]
+            weights_summary = f"\nweighted by {weights_name}"
+        else:
+            weights_summary = "unweighted"
+        group_summary = ", ".join(format_array_flat(self._unique_coord, 30).split())
+        return (
+            f"{self.__class__.__name__}, grouped over {self._unique_coord.name},"
+            f"\n{self._unique_coord.size} groups with labels {group_summary}."
+            f"{weights_summary}"
         )
 
     def _get_index_and_items(self, index, grouper):
@@ -980,7 +1003,20 @@ def reduce(self, func, dim=None, keep_attrs=None, **kwargs):
             keep_attrs = _get_keep_attrs(default=False)
 
         def reduce_dataset(ds):
-            return ds.reduce(func, dim, keep_attrs, **kwargs)
+            if self._weighted:
+                # yuck
+                if func.__name__ in ["sum", "mean"]:
+                    # TODO: fix this
+                    kwargs.pop("numeric_only", None)
+                    return getattr(ds.weighted(ds["__weights__"]), func.__name__)(
+                        dim=dim, keep_attrs=keep_attrs, **kwargs
+                    )
+                else:
+                    raise NotImplementedError(
+                        f"Weighted {func} has not been implemented."
+                    )
+            else:
+                return ds.reduce(func, dim, keep_attrs, **kwargs)
 
         check_reduce_dims(dim, self.dims)