support dask arrays in rolling computations using bottleneck functions

ci/requirements-py36.yml

@@ -16,6 +16,7 @@ dependencies:
   - seaborn
   - toolz
   - rasterio
+  - bottleneck
   - pip:
     - coveralls
     - pytest-cov

doc/installing.rst

@@ -30,7 +30,7 @@ For accelerating xarray
 
 - `bottleneck <https://github.com/kwgoodman/bottleneck>`__: speeds up
   NaN-skipping and rolling window aggregations by a large factor
-  (1.0 or later)
+  (1.1 or later)
 - `cyordereddict <https://github.com/shoyer/cyordereddict>`__: speeds up most
   internal operations with xarray data structures (for python versions < 3.5)
 

doc/whats-new.rst

@@ -32,6 +32,9 @@ Backward Incompatible Changes
 
 - Old numpy < 1.11 and pandas < 0.18 are no longer supported (:issue:`1512`).
   By `Keisuke Fujii <https://github.com/fujiisoup>`_.
+- The minimum supported version bottleneck has increased to 1.1
+  (:issue:`1279`).
+  By `Joe Hamman <https://github.com/jhamman>`_.
 
 Enhancements
 ~~~~~~~~~~~~
@@ -108,6 +111,10 @@ Enhancements
   By `Joe Hamman <https://github.com/jhamman>`_ and
   `Gerrit Holl <https://github.com/gerritholl>`_.
 
+- Support applying rolling window operations using bottleneck's moving window
+  functions on data stored as dask arrays (:issue:`1279`).
+  By `Joe Hamman <https://github.com/jhamman>`_.
+
 Bug fixes
 ~~~~~~~~~
 

xarray/core/dask_array_ops.py

@@ -0,0 +1,26 @@
+"""Define core operations for xarray objects.
+"""
+import numpy as np
+
+try:
+    import dask.array as da
+except ImportError:
+    pass
+
+
+def dask_rolling_wrapper(moving_func, a, window, min_count=None, axis=-1):
+    '''wrapper to apply bottleneck moving window funcs on dask arrays'''
+    # 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)}
+    # create ghosted arrays
+    ag = da.ghost.ghost(a, depth=depth, boundary=boundary)
+    # apply rolling func
+    out = ag.map_blocks(moving_func, window, min_count=min_count,
+                        axis=axis, dtype=a.dtype)
+    # trim array
+    result = da.ghost.trim_internal(out, depth)
+    return result

xarray/core/ops.py

@@ -10,6 +10,7 @@
 from __future__ import print_function
 
 import operator
+import warnings
 from distutils.version import LooseVersion
 
 import numpy as np
@@ -22,6 +23,13 @@
 
 try:
     import bottleneck as bn
+    if LooseVersion(bn.__version__) < LooseVersion('1.1'):
+        warnings.warn(
+            "The installed version of bottleneck {ver} is not supported "
+            "in xarray and will be not be used\nThe minimum supported "
+            "version is 1.1\n".format(ver=bn.__version__), UserWarning)
+        raise ImportError(
+            'bottleneck version {} is not supported'.format(bn.__version__))
     has_bottleneck = True
 except ImportError:
     # use numpy methods instead
@@ -51,7 +59,8 @@
 BOTTLENECK_ROLLING_METHODS = {'move_sum': 'sum', 'move_mean': 'mean',
                               'move_std': 'std', 'move_min': 'min',
                               'move_max': 'max', 'move_var': 'var',
-                              'move_argmin': 'argmin', 'move_argmax': 'argmax'}
+                              'move_argmin': 'argmin', 'move_argmax': 'argmax',
+                              'move_median': 'median'}
 # TODO: wrap take, dot, sort
 
 
@@ -243,6 +252,7 @@ def rolling_count(rolling):
 
     return rolling_count.where(enough_periods)
 
+
 def inject_reduce_methods(cls):
     methods = ([(name, getattr(duck_array_ops, 'array_%s' % name), False)
                 for name in REDUCE_METHODS] +
@@ -354,40 +364,16 @@ def inject_bottleneck_rolling_methods(cls):
     setattr(cls, 'count', func)
 
     # bottleneck rolling methods
-    if has_bottleneck:
-        # TODO: Bump the required version of bottlneck to 1.1 and remove all
-        # these version checks (see GH#1278)
-        bn_version = LooseVersion(bn.__version__)
-        bn_min_version = LooseVersion('1.0')
-        bn_version_1_1 = LooseVersion('1.1')
-        if bn_version < bn_min_version:
-            return
-
-        for bn_name, method_name in BOTTLENECK_ROLLING_METHODS.items():
-            try:
-                f = getattr(bn, bn_name)
-                func = cls._bottleneck_reduce(f)
-                func.__name__ = method_name
-                func.__doc__ = _ROLLING_REDUCE_DOCSTRING_TEMPLATE.format(
-                    name=func.__name__, da_or_ds='DataArray')
-                setattr(cls, method_name, func)
-            except AttributeError as e:
-                # skip functions not in Bottleneck 1.0
-                if ((bn_version < bn_version_1_1) and
-                        (bn_name not in ['move_var', 'move_argmin',
-                                         'move_argmax', 'move_rank'])):
-                    raise e
-
-        # bottleneck rolling methods without min_count (bn.__version__ < 1.1)
-        f = getattr(bn, 'move_median')
-        if bn_version >= bn_version_1_1:
-            func = cls._bottleneck_reduce(f)
-        else:
-            func = cls._bottleneck_reduce_without_min_count(f)
-        func.__name__ = 'median'
+    if not has_bottleneck:
+        return
+
+    for bn_name, method_name in BOTTLENECK_ROLLING_METHODS.items():
+        f = getattr(bn, bn_name)
+        func = cls._bottleneck_reduce(f)
+        func.__name__ = method_name
         func.__doc__ = _ROLLING_REDUCE_DOCSTRING_TEMPLATE.format(
             name=func.__name__, da_or_ds='DataArray')
-        setattr(cls, 'median', func)
+        setattr(cls, method_name, func)
 
 
 def inject_datasetrolling_methods(cls):

xarray/core/rolling.py

@@ -10,6 +10,7 @@
 from .combine import concat
 from .ops import (inject_bottleneck_rolling_methods,
                   inject_datasetrolling_methods, has_bottleneck, bn)
+from .dask_array_ops import dask_rolling_wrapper
 
 
 class Rolling(object):
@@ -75,7 +76,8 @@ def __init__(self, obj, min_periods=None, center=False, **windows):
             self._min_periods = window
         else:
             if min_periods <= 0:
-                raise ValueError('min_periods must be greater than zero or None')
+                raise ValueError(
+                    'min_periods must be greater than zero or None')
 
             self._min_periods = min_periods
         self.center = center
@@ -98,7 +100,6 @@ class DataArrayRolling(Rolling):
     This class adds the following class methods;
     + _reduce_method(cls, func)
     + _bottleneck_reduce(cls, func)
-    + _bottleneck_reduce_without_min_count(cls, func)
 
     These class methods will be used to inject numpy or bottleneck function
     by doing
@@ -237,47 +238,23 @@ def _bottleneck_reduce(cls, func):
         def wrapped_func(self, **kwargs):
             from .dataarray import DataArray
 
-            if isinstance(self.obj.data, dask_array_type):
-                raise NotImplementedError(
-                    'Rolling window operation does not work with dask arrays')
-
             # bottleneck doesn't allow min_count to be 0, although it should
             # work the same as if min_count = 1
             if self.min_periods is not None and self.min_periods == 0:
-                min_count = self.min_periods + 1
+                min_count = 1
             else:
                 min_count = self.min_periods
 
-            values = func(self.obj.data, window=self.window,
-                          min_count=min_count,
-                          axis=self.obj.get_axis_num(self.dim))
-
-            result = DataArray(values, self.obj.coords)
-
-            if self.center:
-                result = self._center_result(result)
-
-            return result
-        return wrapped_func
-
-    @classmethod
-    def _bottleneck_reduce_without_min_count(cls, func):
-        """
-        Methods to return a wrapped function for `media` bottoleneck method.
-        bottleneck's median does not accept min_periods.
-        """
-        def wrapped_func(self, **kwargs):
-            from .dataarray import DataArray
-
-            if self.min_periods is not None:
-                raise ValueError('Rolling.median does not accept min_periods')
+            axis = self.obj.get_axis_num(self.dim)
 
             if isinstance(self.obj.data, dask_array_type):
-                raise NotImplementedError(
-                    'Rolling window operation does not work with dask arrays')
-
-            values = func(self.obj.data, window=self.window,
-                          axis=self.obj.get_axis_num(self.dim))
+                values = dask_rolling_wrapper(func, self.obj.data,
+                                              window=self.window,
+                                              min_count=min_count,
+                                              axis=axis)
+            else:
+                values = func(self.obj.data, window=self.window,
+                              min_count=min_count, axis=axis)
 
             result = DataArray(values, self.obj.coords)
 

xarray/tests/__init__.py

@@ -76,7 +76,7 @@
 
 try:
     import bottleneck
-    if LooseVersion(bottleneck.__version__) < LooseVersion('1.0'):
+    if LooseVersion(bottleneck.__version__) < LooseVersion('1.1'):
         raise ImportError('Fall back to numpy')
     has_bottleneck = True
 except ImportError:

xarray/tests/test_dataarray.py

@@ -2627,6 +2627,18 @@ def da(request):
             [0, np.nan, 1, 2, np.nan, 3, 4, 5, np.nan, 6, 7],
             dims='time')
 
+
+@pytest.fixture
+def da_dask(seed=123):
+    pytest.importorskip('dask.array')
+    rs = np.random.RandomState(seed)
+    times = pd.date_range('2000-01-01', freq='1D', periods=21)
+    values = rs.normal(size=(1, 21, 1))
+    da = DataArray(values, dims=('a', 'time', 'x')).chunk({'time': 7})
+    da['time'] = times
+    return da
+
+
 def test_rolling_iter(da):
 
     rolling_obj = da.rolling(time=7)
@@ -2646,8 +2658,6 @@ def test_rolling_doc(da):
 
 
 def test_rolling_properties(da):
-    pytest.importorskip('bottleneck', minversion='1.0')
-
     rolling_obj = da.rolling(time=4)
 
     assert rolling_obj.obj.get_axis_num('time') == 1
@@ -2669,19 +2679,7 @@ def test_rolling_properties(da):
 @pytest.mark.parametrize('center', (True, False, None))
 @pytest.mark.parametrize('min_periods', (1, None))
 def test_rolling_wrapped_bottleneck(da, name, center, min_periods):
-    pytest.importorskip('bottleneck')
-    import bottleneck as bn
-
-    # skip if median and min_periods bottleneck version < 1.1
-    if ((min_periods == 1) and
-        (name == 'median') and
-            (LooseVersion(bn.__version__) < LooseVersion('1.1'))):
-        pytest.skip('rolling median accepts min_periods for bottleneck 1.1')
-
-    # skip if var and bottleneck version < 1.1
-    if ((name == 'median') and
-            (LooseVersion(bn.__version__) < LooseVersion('1.1'))):
-        pytest.skip('rolling median accepts min_periods for bottleneck 1.1')
+    bn = pytest.importorskip('bottleneck', minversion="1.1")
 
     # Test all bottleneck functions
     rolling_obj = da.rolling(time=7, min_periods=min_periods)
@@ -2698,14 +2696,22 @@ def test_rolling_wrapped_bottleneck(da, name, center, min_periods):
     assert_equal(actual, da['time'])
 
 
-def test_rolling_invalid_args(da):
-    pytest.importorskip('bottleneck', minversion="1.0")
-    import bottleneck as bn
-    if LooseVersion(bn.__version__) >= LooseVersion('1.1'):
-        pytest.skip('rolling median accepts min_periods for bottleneck 1.1')
-    with pytest.raises(ValueError) as exception:
-        da.rolling(time=7, min_periods=1).median()
-    assert 'Rolling.median does not' in str(exception)
+@pytest.mark.parametrize('name', ('sum', 'mean', 'std', 'min', 'max',
+                                  'median'))
+@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.importorskip('dask.array')
+    # dask version
+    rolling_obj = da_dask.rolling(time=7, min_periods=min_periods)
+    actual = getattr(rolling_obj, name)().load()
+    # numpy version
+    rolling_obj = da_dask.load().rolling(time=7, min_periods=min_periods)
+    expected = getattr(rolling_obj, name)()
+
+    # using all-close because rolling over ghost cells introduces some
+    # precision errors
+    assert_allclose(actual, expected)
 
 
 @pytest.mark.parametrize('center', (True, False))
@@ -2718,8 +2724,8 @@ def test_rolling_pandas_compat(da, center, window, min_periods):
     if min_periods is not None and window < min_periods:
         min_periods = window
 
-    s_rolling = pd.rolling_mean(s, window, center=center,
-                                min_periods=min_periods)
+    s_rolling = s.rolling(window, center=center,
+                          min_periods=min_periods).mean()
     da_rolling = da.rolling(index=window, center=center,
                             min_periods=min_periods).mean()
     # pandas does some fancy stuff in the last position,
@@ -2729,6 +2735,7 @@ def test_rolling_pandas_compat(da, center, window, min_periods):
     np.testing.assert_allclose(s_rolling.index,
                                da_rolling['index'])
 
+
 @pytest.mark.parametrize('da', (1, 2), indirect=True)
 @pytest.mark.parametrize('center', (True, False))
 @pytest.mark.parametrize('min_periods', (None, 1, 2, 3))