Add coarsen.construct

Author: dcherian

doc/howdoi.rst

@@ -24,7 +24,7 @@ How do I ...
    * - change the order of dimensions
      - :py:meth:`DataArray.transpose`, :py:meth:`Dataset.transpose`
    * - reshape dimensions
-     - :py:meth:`DataArray.stack`, :py:meth:`Dataset.stack`
+     - :py:meth:`DataArray.stack`, :py:meth:`Dataset.stack`, :py:meth:`Dataset.coarsen.construct`, :py:meth:`DataArray.coarsen.construct`
    * - remove a variable from my object
      - :py:meth:`Dataset.drop_vars`, :py:meth:`DataArray.drop_vars`
    * - remove dimensions of length 1 or 0

doc/whats-new.rst

@@ -21,6 +21,8 @@ v0.18.3 (unreleased)
 
 New Features
 ~~~~~~~~~~~~
+- Added :py:meth:`Dataset.coarsen.construct`, :py:meth:`DataArray.coarsen.construct` (:issue:`5454`, :pull:`5475`).
+  By `Deepak Cherian <https://github.com/dcherian>`_.
 - Allow assigning values to a subset of a dataset using positional or label-based
   indexing (:issue:`3015`, :pull:`5362`).
   By `Matthias Göbel <https://github.com/matzegoebel>`_.

xarray/core/rolling.py

@@ -1,4 +1,5 @@
 import functools
+import itertools
 import warnings
 from typing import Any, Callable, Dict
 
@@ -845,6 +846,80 @@ def __repr__(self):
             klass=self.__class__.__name__, attrs=",".join(attrs)
         )
 
+    def construct(
+        self,
+        window_dim=None,
+        keep_attrs=None,
+        **window_dim_kwargs,
+    ):
+        """
+        Convert this Coarsen object to a Dataset,
+        where the window dimension is reshaped to new dimensions
+
+        Parameters
+        ----------
+        window_dim: str or a mapping, optional
+            A mapping from dimension name to the new window dimension names.
+            Just a string can be used for 1d-rolling.
+        fill_value: optional. Default dtypes.NA
+            Filling value to match the dimension size.
+        **window_dim_kwargs : {dim: new_name, ...}, optional
+            The keyword arguments form of ``window_dim``.
+
+        Returns
+        -------
+        Dataset with variables converted from rolling object.
+        """
+
+        from .dataarray import DataArray
+        from .dataset import Dataset
+
+        if window_dim is None:
+            if len(window_dim_kwargs) == 0:
+                raise ValueError(
+                    "Either window_dim or window_dim_kwargs need to be specified."
+                )
+            window_dim = {d: window_dim_kwargs[d] for d in self.dim}
+
+        if keep_attrs is None:
+            keep_attrs = _get_keep_attrs(default=True)
+
+        missing_dims = set(window_dim) - set(self.windows)
+        if missing_dims:
+            raise ValueError(
+                f"'window_dim' must contain entries for all dimensions to coarsen. Missing {missing_dims}"
+            )
+        missing_windows = set(self.windows) - set(window_dim)
+        if missing_windows:
+            raise ValueError(
+                f"'window_dim' includes dimensions that will not be coarsened: {missing_windows}"
+            )
+
+        reshaped = Dataset()
+        if isinstance(self.obj, DataArray):
+            obj = self.obj._to_temp_dataset()
+        else:
+            obj = self.obj
+
+        for key, var in obj.variables.items():
+            reshaped_dims = tuple(
+                itertools.chain(*[window_dim.get(dim, [dim]) for dim in list(var.dims)])
+            )
+            if reshaped_dims != var.dims:
+                windows = {w: self.windows[w] for w in window_dim if w in var.dims}
+                reshaped_var, _ = var.coarsen_reshape(windows, self.boundary, self.side)
+                attrs = var.attrs if keep_attrs else {}
+                reshaped[key] = (reshaped_dims, reshaped_var, attrs)
+            else:
+                reshaped[key] = var
+
+        should_be_coords = set(window_dim) & set(self.obj.coords)
+        result = reshaped.set_coords(should_be_coords)
+        if isinstance(self.obj, DataArray):
+            return self.obj._from_temp_dataset(result)
+        else:
+            return result
+
 
 class DataArrayCoarsen(Coarsen):
     __slots__ = ()

xarray/core/variable.py

@@ -2158,7 +2158,7 @@ def coarsen(
         if not windows:
             return self._replace(attrs=_attrs)
 
-        reshaped, axes = self._coarsen_reshape(windows, boundary, side)
+        reshaped, axes = self.coarsen_reshape(windows, boundary, side)
         if isinstance(func, str):
             name = func
             func = getattr(duck_array_ops, name, None)
@@ -2167,7 +2167,7 @@ def coarsen(
 
         return self._replace(data=func(reshaped, axis=axes, **kwargs), attrs=_attrs)
 
-    def _coarsen_reshape(self, windows, boundary, side):
+    def coarsen_reshape(self, windows, boundary, side):
         """
         Construct a reshaped-array for coarsen
         """
@@ -2183,7 +2183,9 @@ def _coarsen_reshape(self, windows, boundary, side):
 
         for d, window in windows.items():
             if window <= 0:
-                raise ValueError(f"window must be > 0. Given {window}")
+                raise ValueError(
+                    f"window must be > 0. Given {window} for dimension {d}"
+                )
 
         variable = self
         for d, window in windows.items():
@@ -2193,8 +2195,8 @@ def _coarsen_reshape(self, windows, boundary, side):
             if boundary[d] == "exact":
                 if n * window != size:
                     raise ValueError(
-                        "Could not coarsen a dimension of size {} with "
-                        "window {}".format(size, window)
+                        f"Could not coarsen a dimension of size {size} with "
+                        f"window {window} and boundary='exact'. Try a different 'boundary' option."
                     )
             elif boundary[d] == "trim":
                 if side[d] == "left":

xarray/tests/test_coarsen.py

@@ -5,7 +5,13 @@
 import xarray as xr
 from xarray import DataArray, Dataset, set_options
 
-from . import assert_allclose, assert_equal, has_dask, requires_cftime
+from . import (
+    assert_allclose,
+    assert_equal,
+    has_dask,
+    raise_if_dask_computes,
+    requires_cftime,
+)
 from .test_dataarray import da
 from .test_dataset import ds
 
@@ -299,3 +305,42 @@ def test_coarsen_da_reduce(da, window, name):
     actual = coarsen_obj.reduce(getattr(np, f"nan{name}"))
     expected = getattr(coarsen_obj, name)()
     assert_allclose(actual, expected)
+
+
+@pytest.mark.parametrize("dask", [True, False])
+def test_coarsen_construct(dask):
+
+    ds = Dataset(
+        {
+            "vart": ("time", np.arange(48)),
+            "varx": ("x", np.arange(10)),
+            "vartx": (("x", "time"), np.arange(480).reshape(10, 48)),
+            "vary": ("y", np.arange(12)),
+        },
+        coords={"time": np.arange(48), "y": np.arange(12)},
+    )
+
+    if dask and has_dask:
+        ds = ds.chunk({"x": 4, "time": 10})
+    with raise_if_dask_computes():
+        actual = ds.coarsen(time=12, x=5).construct(
+            {"time": ("year", "month"), "x": ("x", "x_reshaped")}
+        )
+
+    expected = xr.Dataset()
+    expected["vart"] = (("year", "month"), ds.vart.data.reshape((-1, 12)))
+    expected["varx"] = (("x", "x_reshaped"), ds.varx.data.reshape((-1, 5)))
+    expected["vartx"] = (
+        ("x", "x_reshaped", "year", "month"),
+        ds.vartx.data.reshape(2, 5, 4, 12),
+    )
+    expected["vary"] = ds.vary
+    expected.coords["time"] = (("year", "month"), ds.time.data.reshape((-1, 12)))
+
+    assert_equal(actual, expected)
+
+    with raise_if_dask_computes():
+        actual = ds.vartx.coarsen(time=12, x=5).construct(
+            {"time": ("year", "month"), "x": ("x", "x_reshaped")}
+        )
+    assert_equal(actual, expected["vartx"])