Remove xarray_groupby_reduce

Author: dcherian

flox/xarray.py

@@ -16,7 +16,7 @@
 from .xrutils import is_duck_dask_array, isnull
 
 if TYPE_CHECKING:
-    from xarray import DataArray, Dataset, GroupBy, Resample
+    from xarray import DataArray, Dataset, Resample
 
 
 def _get_input_core_dims(group_names, dim, ds, to_group):
@@ -403,47 +403,6 @@ def wrapper(array, to_group, *, func, skipna, **kwargs):
         return actual
 
 
-def xarray_groupby_reduce(
-    groupby: "GroupBy",
-    func: Union[str, Aggregation],
-    split_out: int = 1,
-    method: str = "map-reduce",
-    keep_attrs: bool = True,
-):
-    """Apply on an existing Xarray groupby object for convenience."""
-
-    def wrapper(*args, **kwargs):
-        result, _ = groupby_reduce(*args, **kwargs)
-        return result
-
-    groups = list(groupby.groups.keys())
-    outdim = groupby._unique_coord.name
-    groupdim = groupby._group_dim
-
-    actual = xr.apply_ufunc(
-        wrapper,
-        groupby._obj,
-        groupby._group,
-        input_core_dims=[[groupdim], [groupdim]],
-        # for xarray's test_groupby_duplicate_coordinate_labels
-        exclude_dims=set(groupdim),
-        output_core_dims=[[outdim]],
-        dask="allowed",
-        dask_gufunc_kwargs=dict(output_sizes={outdim: len(groups)}),
-        keep_attrs=keep_attrs,
-        kwargs={
-            "func": func,
-            "axis": -1,
-            "split_out": split_out,
-            "expected_groups": groups,
-            "method": method,
-        },
-    )
-    actual[outdim] = groups
-
-    return actual
-
-
 def rechunk_for_cohorts(
     obj: Union["DataArray", "Dataset"],
     dim: str,

tests/test_xarray.py

@@ -6,12 +6,7 @@
 xr = pytest.importorskip("xarray")
 # isort: on
 
-from flox.xarray import (
-    rechunk_to_group_boundaries,
-    resample_reduce,
-    xarray_groupby_reduce,
-    xarray_reduce,
-)
+from flox.xarray import rechunk_to_group_boundaries, resample_reduce, xarray_reduce
 
 from . import assert_equal, engine, has_dask, raise_if_dask_computes, requires_dask
 
@@ -71,32 +66,6 @@ def test_xarray_reduce(skipna, add_nan, min_count, engine):
     # assert_equal(expected, actual)
 
 
-def test_xarray_groupby_reduce():
-    arr = np.ones((4, 12))
-
-    labels = np.array(["a", "a", "c", "c", "c", "b", "b", "c", "c", "b", "b", "f"])
-    labels = np.array(labels)
-    labels2 = np.array([1, 2, 2, 1])
-
-    da = xr.DataArray(
-        arr, dims=("x", "y"), coords={"labels2": ("x", labels2), "labels": ("y", labels)}
-    ).expand_dims(z=4)
-
-    grouped = da.groupby("labels")
-    expected = grouped.mean()
-    actual = xarray_groupby_reduce(grouped, "mean")
-    assert_equal(expected, actual)
-
-    actual = xarray_groupby_reduce(da.transpose("y", ...).groupby("labels"), "mean")
-    assert_equal(expected, actual)
-
-    # TODO: fails because of stacking
-    # grouped = da.groupby("labels2")
-    # expected = grouped.mean()
-    # actual = xarray_groupby_reduce(grouped, "mean")
-    # assert_equal(expected, actual)
-
-
 def test_xarray_reduce_multiple_groupers():
     arr = np.ones((4, 12))