boundary options for rolling.construct

ci/requirements/py36-min-all-deps.yml

@@ -15,8 +15,8 @@ dependencies:
   - cfgrib=0.9
   - cftime=1.0
   - coveralls
-  - dask=1.2
-  - distributed=1.27
+  - dask=1.2.2
+  - distributed=1.28
   - flake8
   - h5netcdf=0.7
   - h5py=2.9  # Policy allows for 2.10, but it's a conflict-fest

doc/whats-new.rst

@@ -25,6 +25,10 @@ Breaking changes
 
 New Features
 ~~~~~~~~~~~~
+- Some boundary options ('edge' | 'reflect' | 'symmetric' | 'wrap') are 
+  newly supported in :py:meth:`~xarray.DataArray.rolling.construct`.
+  (:issue:`2007`)
+  By `Keisuke Fujii <https://github.com/fujiisoup>`_.
 - :py:meth:`Dataset.quantile`, :py:meth:`DataArray.quantile` and ``GroupBy.quantile``
   now work with dask Variables.
   By `Deepak Cherian <https://github.com/dcherian>`_.

xarray/core/dask_array_ops.py

@@ -27,7 +27,7 @@ def dask_rolling_wrapper(moving_func, a, window, min_count=None, axis=-1):
     return result
 
 
-def rolling_window(a, axis, window, center, fill_value):
+def rolling_window(a, axis, window, center, fill_value, mode):
     """Dask's equivalence to np.utils.rolling_window
     """
     import dask.array as da
@@ -68,25 +68,44 @@ def rolling_window(a, axis, window, center, fill_value):
             # smaller than the depth, we enlarge this and truncate it later.
             drop_size = depth[axis] - pad_size
             pad_size = depth[axis]
-        shape = list(a.shape)
-        shape[axis] = pad_size
-        chunks = list(a.chunks)
-        chunks[axis] = (pad_size,)
-        fill_array = da.full(shape, fill_value, dtype=a.dtype, chunks=chunks)
-        a = da.concatenate([fill_array, a], axis=axis)
 
-    boundary = {d: fill_value for d in range(a.ndim)}
+    # pad manually
+    pads = (
+        ((0, 0),) * axis
+        + ((pad_size + depth[axis], depth[axis]),)
+        + ((0, 0),) * (a.ndim - axis - 1)
+    )
+    if mode is None:
+        a = da.pad(a, pads, mode="constant", constant_values=fill_value)
+    else:
+        a = da.pad(a, pads, mode=mode)
+
+    # merge the padded part into the edge block
+    new_chunks = list(a.chunks[axis])
+    first_chunk = new_chunks.pop(0)
+    last_chunk = new_chunks.pop(-1)
+    if len(new_chunks) == 0:
+        new_chunks = (first_chunk + last_chunk,)
+    else:
+        new_chunks[0] += first_chunk
+        new_chunks[-1] += last_chunk
 
-    # create overlap arrays
+    chunks = list(a.chunks)
+    chunks[axis] = tuple(new_chunks)
+    a = da.rechunk(a, tuple(chunks))
+
+    # create overlap arrays without boundary
+    boundary = {d: "none" for d in range(a.ndim)}
     ag = da.overlap.overlap(a, depth=depth, boundary=boundary)
 
-    # apply rolling func
     def func(x, window, axis=-1):
         x = np.asarray(x)
         rolling = nputils._rolling_window(x, window, axis)
         return rolling[(slice(None),) * axis + (slice(offset, None),)]
 
-    chunks = list(a.chunks)
+    # we need a correct shape
+    chunks = list(ag.chunks)
+    chunks[axis] = tuple([c - (window - 1 + offset) for c in chunks[axis]])
     chunks.append(window)
     out = ag.map_blocks(
         func, dtype=a.dtype, new_axis=a.ndim, chunks=chunks, window=window, axis=axis

xarray/core/duck_array_ops.py

@@ -501,12 +501,14 @@ def last(values, axis, skipna=None):
     return take(values, -1, axis=axis)
 
 
-def rolling_window(array, axis, window, center, fill_value):
+def rolling_window(array, axis, window, center, fill_value, mode):
     """
     Make an ndarray with a rolling window of axis-th dimension.
     The rolling dimension will be placed at the last dimension.
     """
     if isinstance(array, dask_array_type):
-        return dask_array_ops.rolling_window(array, axis, window, center, fill_value)
+        return dask_array_ops.rolling_window(
+            array, axis, window, center, fill_value, mode
+        )
     else:  # np.ndarray
-        return nputils.rolling_window(array, axis, window, center, fill_value)
+        return nputils.rolling_window(array, axis, window, center, fill_value, mode)

xarray/core/nputils.py

@@ -140,7 +140,7 @@ def __setitem__(self, key, value):
         self._array[key] = np.moveaxis(value, vindex_positions, mixed_positions)
 
 
-def rolling_window(a, axis, window, center, fill_value):
+def rolling_window(a, axis, window, center, fill_value, mode):
     """ rolling window with padding. """
     pads = [(0, 0) for s in a.shape]
     if center:
@@ -149,7 +149,10 @@ def rolling_window(a, axis, window, center, fill_value):
         pads[axis] = (start, end)
     else:
         pads[axis] = (window - 1, 0)
-    a = np.pad(a, pads, mode="constant", constant_values=fill_value)
+    if mode is None:
+        a = np.pad(a, pads, mode="constant", constant_values=fill_value)
+    else:
+        a = np.pad(a, pads, mode=mode)
     return _rolling_window(a, window, axis)
 
 

xarray/core/rolling.py

@@ -193,7 +193,7 @@ def __iter__(self):
 
             yield (label, window)
 
-    def construct(self, window_dim, stride=1, fill_value=dtypes.NA):
+    def construct(self, window_dim, stride=1, fill_value=dtypes.NA, mode=None):
         """
         Convert this rolling object to xr.DataArray,
         where the window dimension is stacked as a new dimension
@@ -206,6 +206,11 @@ def construct(self, window_dim, stride=1, fill_value=dtypes.NA):
             Size of stride for the rolling window.
         fill_value: optional. Default dtypes.NA
             Filling value to match the dimension size.
+        mode: optional. Default None
+            One of None | 'edge' | 'reflect' | 'symmetric' | 'wrap'
+            For the details of the mode, see
+            https://docs.scipy.org/doc/numpy/reference/generated/numpy.pad.html
+            If it is not None, fill_value is ignored.
 
         Returns
         -------
@@ -234,7 +239,12 @@ def construct(self, window_dim, stride=1, fill_value=dtypes.NA):
         from .dataarray import DataArray
 
         window = self.obj.variable.rolling_window(
-            self.dim, self.window, window_dim, self.center, fill_value=fill_value
+            self.dim,
+            self.window,
+            window_dim,
+            self.center,
+            fill_value=fill_value,
+            mode=mode,
         )
         result = DataArray(
             window, dims=self.obj.dims + (window_dim,), coords=self.obj.coords
@@ -466,7 +476,7 @@ def _numpy_or_bottleneck_reduce(
             **kwargs,
         )
 
-    def construct(self, window_dim, stride=1, fill_value=dtypes.NA):
+    def construct(self, window_dim, stride=1, fill_value=dtypes.NA, mode=None):
         """
         Convert this rolling object to xr.Dataset,
         where the window dimension is stacked as a new dimension
@@ -479,6 +489,10 @@ def construct(self, window_dim, stride=1, fill_value=dtypes.NA):
             size of stride for the rolling window.
         fill_value: optional. Default dtypes.NA
             Filling value to match the dimension size.
+        mode: optional. Default None
+            One of None | 'edge' | 'reflect' | 'symmetric' | 'wrap'
+            For the details of the mode, see
+            https://docs.scipy.org/doc/numpy/reference/generated/numpy.pad.html
 
         Returns
         -------
@@ -491,7 +505,7 @@ def construct(self, window_dim, stride=1, fill_value=dtypes.NA):
         for key, da in self.obj.data_vars.items():
             if self.dim in da.dims:
                 dataset[key] = self.rollings[key].construct(
-                    window_dim, fill_value=fill_value
+                    window_dim, fill_value=fill_value, mode=mode
                 )
             else:
                 dataset[key] = da

xarray/core/variable.py

@@ -1806,7 +1806,7 @@ def rank(self, dim, pct=False):
         return Variable(self.dims, ranked)
 
     def rolling_window(
-        self, dim, window, window_dim, center=False, fill_value=dtypes.NA
+        self, dim, window, window_dim, center=False, fill_value=dtypes.NA, mode=None
     ):
         """
         Make a rolling_window along dim and add a new_dim to the last place.
@@ -1824,6 +1824,9 @@ def rolling_window(
             of the axis.
         fill_value:
             value to be filled.
+        mode: optional. Default None
+            One of None | 'edge' | 'reflect' | 'symmetric' | 'wrap'
+            If it is not None, fill_value is ignored.
 
         Returns
         -------
@@ -1861,6 +1864,7 @@ def rolling_window(
                 window=window,
                 center=center,
                 fill_value=fill_value,
+                mode=mode,
             ),
         )
 

xarray/tests/test_dataarray.py

@@ -4257,6 +4257,14 @@ def test_rolling_wrapped_dask_nochunk(center):
     assert_allclose(actual, expected)
 
 
+@pytest.mark.parametrize("da", (1, 2), indirect=True)
+@pytest.mark.parametrize("center", [True, False])
+@pytest.mark.parametrize("mode", ["edge", "symmetric", "reflect", "wrap"])
+def test_rolling_mode(da, center, mode):
+    # just a working test
+    da.rolling(time=7, center=center).construct("constructed", mode=mode)
+
+
 @pytest.mark.parametrize("center", (True, False))
 @pytest.mark.parametrize("min_periods", (None, 1, 2, 3))
 @pytest.mark.parametrize("window", (1, 2, 3, 4))

xarray/tests/test_duck_array_ops.py

@@ -503,10 +503,10 @@ def test_dask_rolling(axis, window, center):
     dx = da.from_array(x, chunks=[(6, 30, 30, 20, 14), 8])
 
     expected = rolling_window(
-        x, axis=axis, window=window, center=center, fill_value=np.nan
+        x, axis=axis, window=window, center=center, fill_value=np.nan, mode=None
     )
     actual = rolling_window(
-        dx, axis=axis, window=window, center=center, fill_value=np.nan
+        dx, axis=axis, window=window, center=center, fill_value=np.nan, mode=None
     )
     assert isinstance(actual, da.Array)
     assert_array_equal(actual, expected)
@@ -515,7 +515,9 @@ def test_dask_rolling(axis, window, center):
     # we need to take care of window size if chunk size is small
     # window/2 should be smaller than the smallest chunk size.
     with pytest.raises(ValueError):
-        rolling_window(dx, axis=axis, window=100, center=center, fill_value=np.nan)
+        rolling_window(
+            dx, axis=axis, window=100, center=center, fill_value=np.nan, mode=None
+        )
 
 
 @pytest.mark.skipif(not has_dask, reason="This is for dask.")

xarray/tests/test_nputils.py

@@ -33,17 +33,35 @@ def test_vindex():
 def test_rolling():
     x = np.array([1, 2, 3, 4], dtype=float)
 
-    actual = rolling_window(x, axis=-1, window=3, center=True, fill_value=np.nan)
+    actual = rolling_window(
+        x, axis=-1, window=3, center=True, fill_value=np.nan, mode=None
+    )
     expected = np.array(
         [[np.nan, 1, 2], [1, 2, 3], [2, 3, 4], [3, 4, np.nan]], dtype=float
     )
     assert_array_equal(actual, expected)
 
-    actual = rolling_window(x, axis=-1, window=3, center=False, fill_value=0.0)
+    actual = rolling_window(
+        x, axis=-1, window=3, center=False, fill_value=0.0, mode=None
+    )
     expected = np.array([[0, 0, 1], [0, 1, 2], [1, 2, 3], [2, 3, 4]], dtype=float)
     assert_array_equal(actual, expected)
 
+    actual = rolling_window(
+        x, axis=-1, window=3, center=True, fill_value=None, mode="wrap"
+    )
+    expected = np.array([[4, 1, 2], [1, 2, 3], [2, 3, 4], [3, 4, 1]], dtype=float)
+    assert_array_equal(actual, expected)
+
+    actual = rolling_window(
+        x, axis=-1, window=3, center=True, fill_value=None, mode="reflect"
+    )
+    expected = np.array([[2, 1, 2], [1, 2, 3], [2, 3, 4], [3, 4, 3]], dtype=float)
+    assert_array_equal(actual, expected)
+
     x = np.stack([x, x * 1.1])
-    actual = rolling_window(x, axis=-1, window=3, center=False, fill_value=0.0)
+    actual = rolling_window(
+        x, axis=-1, window=3, center=True, fill_value=None, mode="reflect"
+    )
     expected = np.stack([expected, expected * 1.1], axis=0)
     assert_array_equal(actual, expected)

xarray/tests/test_variable.py

@@ -821,6 +821,10 @@ def test_pad(self):
             assert_array_equal(actual, expected)
 
     def test_rolling_window(self):
+        v = self.cls(["x", "y", "z"], np.arange(40 * 30 * 2).reshape(40, 30, 2))
+        self._test_rolling_window(v)
+
+    def _test_rolling_window(self, v):
         # Just a working test. See test_nputils for the algorithm validation
         v = self.cls(["x", "y", "z"], np.arange(40 * 30 * 2).reshape(40, 30, 2))
         for (d, w) in [("x", 3), ("y", 5)]:
@@ -836,6 +840,15 @@ def test_rolling_window(self):
             v_loaded = v.load().rolling_window(d, w, d + "_window", center=True)
             assert_array_equal(v_rolling, v_loaded)
 
+            # dask and numpy result should be the same
+            v_rolling = v.rolling_window(
+                d, w, d + "_window", center=True, mode="symmetric"
+            )
+            v_loaded = v.load().rolling_window(
+                d, w, d + "_window", center=True, mode="symmetric"
+            )
+            assert_array_equal(v_rolling, v_loaded)
+
             # numpy backend should not be over-written
             if isinstance(v._data, np.ndarray):
                 with pytest.raises(ValueError):
@@ -1871,6 +1884,11 @@ def test_getitem_with_mask_nd_indexer(self):
             self.cls(("x", "y"), [[0, -1], [-1, 2]]),
         )
 
+    def test_rolling_window(self):
+        v = self.cls(["x", "y", "z"], np.arange(40 * 30 * 2).reshape(40, 30, 2))
+        self._test_rolling_window(v)
+        self._test_rolling_window(v.chunk({"x": 4, "y": 5, "z": -1}))
+
 
 @requires_sparse
 class TestVariableWithSparse: