Fixes centerized rolling with bottleneck. Also, fixed rolling with an integer dask array. (#2122)

Author: fujiisoup

doc/whats-new.rst

@@ -49,6 +49,9 @@ Enhancements
 Bug fixes
 ~~~~~~~~~
 
+- Fixed a bug in `rolling` with bottleneck. Also, fixed a bug in rolling an
+  integer dask array. (:issue:`21133`)
+  By `Keisuke Fujii <https://github.com/fujiisoup>`_.
 - Fixed a bug where `keep_attrs=True` flag was neglected if
   :py:func:`apply_func` was used with :py:class:`Variable`. (:issue:`2114`)
   By `Keisuke Fujii <https://github.com/fujiisoup>`_.

xarray/core/dask_array_ops.py

@@ -3,6 +3,7 @@
 import numpy as np
 
 from . import nputils
+from . import dtypes
 
 try:
     import dask.array as da
@@ -12,12 +13,14 @@
 
 def dask_rolling_wrapper(moving_func, a, window, min_count=None, axis=-1):
     '''wrapper to apply bottleneck moving window funcs on dask arrays'''
+    dtype, fill_value = dtypes.maybe_promote(a.dtype)
+    a = a.astype(dtype)
     # inputs for ghost
     if axis < 0:
         axis = a.ndim + axis
     depth = {d: 0 for d in range(a.ndim)}
     depth[axis] = window - 1
-    boundary = {d: np.nan for d in range(a.ndim)}
+    boundary = {d: fill_value for d in range(a.ndim)}
     # create ghosted arrays
     ag = da.ghost.ghost(a, depth=depth, boundary=boundary)
     # apply rolling func

xarray/core/rolling.py

@@ -285,18 +285,26 @@ def wrapped_func(self, **kwargs):
 
             padded = self.obj.variable
             if self.center:
-                shift = (-self.window // 2) + 1
-
                 if (LooseVersion(np.__version__) < LooseVersion('1.13') and
                         self.obj.dtype.kind == 'b'):
                     # with numpy < 1.13 bottleneck cannot handle np.nan-Boolean
                     # mixed array correctly. We cast boolean array to float.
                     padded = padded.astype(float)
+
+                if isinstance(padded.data, dask_array_type):
+                    # Workaround to make the padded chunk size is larger than
+                    # self.window-1
+                    shift = - (self.window - 1)
+                    offset = -shift - self.window // 2
+                    valid = (slice(None), ) * axis + (
+                        slice(offset, offset + self.obj.shape[axis]), )
+                else:
+                    shift = (-self.window // 2) + 1
+                    valid = (slice(None), ) * axis + (slice(-shift, None), )
                 padded = padded.pad_with_fill_value(**{self.dim: (0, -shift)})
-                valid = (slice(None), ) * axis + (slice(-shift, None), )
 
             if isinstance(padded.data, dask_array_type):
-                values = dask_rolling_wrapper(func, self.obj.data,
+                values = dask_rolling_wrapper(func, padded,
                                               window=self.window,
                                               min_count=min_count,
                                               axis=axis)

xarray/tests/test_dataarray.py

@@ -3439,23 +3439,43 @@ def test_rolling_wrapped_bottleneck(da, name, center, min_periods):
     assert_equal(actual, da['time'])
 
 
-@pytest.mark.parametrize('name', ('sum', 'mean', 'std', 'min', 'max',
-                                  'median'))
+@pytest.mark.parametrize('name', ('mean', 'count'))
 @pytest.mark.parametrize('center', (True, False, None))
 @pytest.mark.parametrize('min_periods', (1, None))
-def test_rolling_wrapped_bottleneck_dask(da_dask, name, center, min_periods):
+@pytest.mark.parametrize('window', (7, 8))
+def test_rolling_wrapped_dask(da_dask, name, center, min_periods, window):
     pytest.importorskip('dask.array')
     # dask version
-    rolling_obj = da_dask.rolling(time=7, min_periods=min_periods)
+    rolling_obj = da_dask.rolling(time=window, min_periods=min_periods,
+                                  center=center)
     actual = getattr(rolling_obj, name)().load()
     # numpy version
-    rolling_obj = da_dask.load().rolling(time=7, min_periods=min_periods)
+    rolling_obj = da_dask.load().rolling(time=window, min_periods=min_periods,
+                                         center=center)
     expected = getattr(rolling_obj, name)()
 
     # using all-close because rolling over ghost cells introduces some
     # precision errors
     assert_allclose(actual, expected)
 
+    # with zero chunked array GH:2113
+    rolling_obj = da_dask.chunk().rolling(time=window, min_periods=min_periods,
+                                          center=center)
+    actual = getattr(rolling_obj, name)().load()
+    assert_allclose(actual, expected)
+
+
+@pytest.mark.parametrize('center', (True, None))
+def test_rolling_wrapped_dask_nochunk(center):
+    # GH:2113
+    pytest.importorskip('dask.array')
+
+    da_day_clim = xr.DataArray(np.arange(1, 367),
+                               coords=[np.arange(1, 367)], dims='dayofyear')
+    expected = da_day_clim.rolling(dayofyear=31, center=center).mean()
+    actual = da_day_clim.chunk().rolling(dayofyear=31, center=center).mean()
+    assert_allclose(actual, expected)
+
 
 @pytest.mark.parametrize('center', (True, False))
 @pytest.mark.parametrize('min_periods', (None, 1, 2, 3))