Rolling window with as_strided

doc/computation.rst

@@ -158,20 +158,31 @@ Aggregation and summary methods can be applied directly to the ``Rolling`` objec
     r.mean()
     r.reduce(np.std)
 
-Note that rolling window aggregations are much faster (both asymptotically and
-because they avoid a loop in Python) when bottleneck_ is installed. Otherwise,
-we fall back to a slower, pure Python implementation.
+Note that rolling window aggregations are faster when bottleneck_ is installed.
 
 .. _bottleneck: https://github.com/kwgoodman/bottleneck/
 
-Finally, we can manually iterate through ``Rolling`` objects:
+We can also manually iterate through ``Rolling`` objects:
 
 .. ipython:: python
 
    @verbatim
    for label, arr_window in r:
       # arr_window is a view of x
 
+Finally, the rolling object has ``to_dataarray`` method, which gives a
+view of the original ``DataArray`` with the windowed dimension is attached to
+the last position.
+You can use this for more advanced rolling operations, such as strided rolling,
+windowed rolling, convolution and short-time FFT.
+
+.. ipython:: python
+
+    rolling_da = r.to_dataarray('window_dim')
+    rolling_da
+    # rolling mean for every 2 points
+    rolling_da.isel(y=slice(None, None, 2)).mean('window_dim')
+
 .. _compute.broadcasting:
 
 Broadcasting by dimension name

doc/whats-new.rst

@@ -27,6 +27,13 @@ Documentation
 
 Enhancements
 ~~~~~~~~~~~~
+- Improve :py:func:`~xarray.DataArray.rooling` logic for speed up.
+  :py:func:`~xarray.DataArrayRolling` object now support ``to_dataarray``
+  method that returns a view of the DataArray object with the rolling-window
+  dimension added to the last position. This enables more flexible operation,
+  such as strided rolling, windowed rolling, ND-rolling, and convolution.
+  (:issue:`1831`, :issue:`1142`, :issue:`819`)
+  By `Keisuke Fujii <https://github.com/fujiisoup>`_.
 - Added nodatavals attribute to DataArray when using :py:func:`~xarray.open_rasterio`. (:issue:`1736`).
   By `Alan Snow <https://github.com/snowman2>`_.
 

xarray/core/dtypes.py

@@ -34,6 +34,8 @@ def maybe_promote(dtype):
         fill_value = np.datetime64('NaT')
     elif np.issubdtype(dtype, np.timedelta64):
         fill_value = np.timedelta64('NaT')
+    elif dtype.kind == 'b':
+        fill_value = False
     else:
         dtype = object
         fill_value = np.nan

xarray/core/duck_array_ops.py

@@ -16,6 +16,7 @@
 import pandas as pd
 
 from . import npcompat
+from . import nputils
 from . import dtypes
 from .pycompat import dask_array_type
 from .nputils import nanfirst, nanlast
@@ -263,3 +264,24 @@ def last(values, axis, skipna=None):
         _fail_on_dask_array_input_skipna(values)
         return nanlast(values, axis)
     return take(values, -1, axis=axis)
+
+
+def rolling_window(array, axis, window):
+    """
+    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):
+        if window < 1:
+            raise ValueError(
+                "`window` must be at least 1. Given : {}".format(window))
+        if window > array.shape[axis]:
+            raise ValueError("`window` is too long. Given : {}".format(window))
+
+        axis = nputils._validate_axis(array, axis)
+        size = array.shape[axis] - window + 1
+        arrays = [array[(slice(None), ) * axis + (slice(w, size + w), )]
+                  for w in range(window)]
+        return da.stack(arrays, axis=-1)
+    else:  # np.ndarray
+        return nputils.rolling_window(array, axis, window)

xarray/core/npcompat.py

@@ -2,6 +2,18 @@
 from __future__ import division
 from __future__ import print_function
 import numpy as np
+from distutils.version import LooseVersion
+
+
+if LooseVersion(np.__version__) < LooseVersion('1.12'):
+    def as_strided(x, shape=None, strides=None, subok=False, writeable=True):
+        array = np.lib.stride_tricks.as_strided(x, shape, strides, subok)
+        array.setflags(write=writeable)
+        return array
+
+else:
+    as_strided = np.lib.stride_tricks.as_strided
+
 
 try:
     from numpy import nancumsum, nancumprod, flip

xarray/core/nputils.py

@@ -4,6 +4,7 @@
 import numpy as np
 import pandas as pd
 import warnings
+from . import npcompat
 
 
 def _validate_axis(data, axis):
@@ -133,3 +134,52 @@ def __setitem__(self, key, value):
         mixed_positions, vindex_positions = _advanced_indexer_subspaces(key)
         self._array[key] = np.moveaxis(value, vindex_positions,
                                        mixed_positions)
+
+
+def rolling_window(a, axis, window):
+    """
+    Make an ndarray with a rolling window along axis.
+
+    Parameters
+    ----------
+    a : array_like
+        Array to add rolling window to
+    axis: int
+        axis position along which rolling window will be applied.
+    window : int
+        Size of rolling window
+
+    Returns
+    -------
+    Array that is a view of the original array with a added dimension
+    of size w.
+
+    Examples
+    --------
+    >>> x=np.arange(10).reshape((2,5))
+    >>> np.rolling_window(x, 3, axis=-1)
+    array([[[0, 1, 2], [1, 2, 3], [2, 3, 4]],
+           [[5, 6, 7], [6, 7, 8], [7, 8, 9]]])
+
+    Calculate rolling mean of last dimension:
+    >>> np.mean(np.rolling_window(x, 3, axis=-1), -1)
+    array([[ 1.,  2.,  3.],
+           [ 6.,  7.,  8.]])
+
+    This function is taken from https://github.com/numpy/numpy/pull/31
+    but slightly modified to accept axis option.
+    """
+    axis = _validate_axis(a, axis)
+    a = np.swapaxes(a, axis, -1)
+
+    if window < 1:
+        raise ValueError(
+            "`window` must be at least 1. Given : {}".format(window))
+    if window > a.shape[-1]:
+        raise ValueError("`window` is too long. Given : {}".format(window))
+
+    shape = a.shape[:-1] + (a.shape[-1] - window + 1, window)
+    strides = a.strides + (a.strides[-1],)
+    rolling = npcompat.as_strided(a, shape=shape, strides=strides,
+                                  writeable=False)
+    return np.swapaxes(rolling, -2, axis)

xarray/core/rolling.py

@@ -6,8 +6,6 @@
 from distutils.version import LooseVersion
 
 from .pycompat import OrderedDict, zip, dask_array_type
-from .common import full_like
-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
@@ -127,62 +125,75 @@ class DataArrayRolling(Rolling):
     def __init__(self, obj, min_periods=None, center=False, **windows):
         super(DataArrayRolling, self).__init__(obj, min_periods=min_periods,
                                                center=center, **windows)
-        self._windows = None
-        self._valid_windows = None
         self.window_indices = None
         self.window_labels = None
 
         self._setup_windows()
 
-    @property
-    def windows(self):
-        if self._windows is None:
-            self._windows = OrderedDict(zip(self.window_labels,
-                                            self.window_indices))
-        return self._windows
-
     def __iter__(self):
-        for (label, indices, valid) in zip(self.window_labels,
-                                           self.window_indices,
-                                           self._valid_windows):
-
+        for (label, indices) in zip(self.window_labels, self.window_indices):
             window = self.obj.isel(**{self.dim: indices})
 
-            if not valid:
-                window = full_like(window, fill_value=True, dtype=bool)
+            counts = window.count(dim=self.dim)
+            window = window.where(counts >= self._min_periods)
 
             yield (label, window)
 
     def _setup_windows(self):
         """
         Find the indices and labels for each window
         """
-        from .dataarray import DataArray
-
         self.window_labels = self.obj[self.dim]
-
         window = int(self.window)
-
         dim_size = self.obj[self.dim].size
 
         stops = np.arange(dim_size) + 1
         starts = np.maximum(stops - window, 0)
 
-        if self._min_periods > 1:
-            valid_windows = (stops - starts) >= self._min_periods
-        else:
-            # No invalid windows
-            valid_windows = np.ones(dim_size, dtype=bool)
-        self._valid_windows = DataArray(valid_windows, dims=(self.dim, ),
-                                        coords=self.obj[self.dim].coords)
-
         self.window_indices = [slice(start, stop)
                                for start, stop in zip(starts, stops)]
 
-    def _center_result(self, result):
-        """center result"""
-        shift = (-self.window // 2) + 1
-        return result.shift(**{self.dim: shift})
+    def to_dataarray(self, window_dim):
+        """
+        Convert this rolling object to xr.DataArray,
+        where the window dimension is stacked as a new dimension
+
+        Parameters
+        ----------
+        window_dim: str
+            New name of the window dimension.
+
+        Returns
+        -------
+        DataArray that is a view of the original array.
+
+        Note
+        ----
+        The return array is not writeable.
+
+        Examples
+        --------
+        >>> da = DataArray(np.arange(8).reshape(2, 4), dims=('a', 'b'))
+
+        >>> da.rolling_window(x, 'b', 4, 'window_dim')
+        <xarray.DataArray (a: 2, b: 4, window_dim: 3)>
+        array([[[np.nan, np.nan, 0], [np.nan, 0, 1], [0, 1, 2], [1, 2, 3]],
+               [[np.nan, np.nan, 4], [np.nan, 4, 5], [4, 5, 6], [5, 6, 7]]])
+        Dimensions without coordinates: a, b, window_dim
+
+        >>> da.rolling_window(x, 'b', 4, 'window_dim', center=True)
+        <xarray.DataArray (a: 2, b: 4, window_dim: 3)>
+        array([[[np.nan, 0, 1], [0, 1, 2], [1, 2, 3], [2, 3, np.nan]],
+               [[np.nan, 4, 5], [4, 5, 6], [5, 6, 7], [6, 7, np.nan]]])
+        Dimensions without coordinates: a, b, window_dim
+        """
+
+        from .dataarray import DataArray
+
+        window = self.obj.variable.rolling_window(self.dim, self.window,
+                                                  window_dim, self.center)
+        return DataArray(window, dims=self.obj.dims + (window_dim,),
+                         coords=self.obj.coords)
 
     def reduce(self, func, **kwargs):
         """Reduce the items in this group by applying `func` along some
@@ -203,26 +214,18 @@ def reduce(self, func, **kwargs):
             Array with summarized data.
         """
 
-        windows = [window.reduce(func, dim=self.dim, **kwargs)
-                   for _, window in self]
-
-        # Find valid windows based on count
-        if self.dim in self.obj.coords:
-            concat_dim = self.window_labels
-        else:
-            concat_dim = self.dim
-        counts = concat([window.count(dim=self.dim) for _, window in self],
-                        dim=concat_dim)
-        result = concat(windows, dim=concat_dim)
-        # restore dim order
-        result = result.transpose(*self.obj.dims)
+        windows = self.to_dataarray('_rolling_window_dim')
+        result = windows.reduce(func, dim='_rolling_window_dim', **kwargs)
 
+        # Find valid windows based on count.
+        # We do not use `reduced.count()` because it constructs a larger array
+        # (notice that `windows` is just a view)
+        counts = (~self.obj.isnull()).rolling(
+            center=self.center, **{self.dim: self.window}).to_dataarray(
+                '_rolling_window_dim').sum(dim='_rolling_window_dim')
         result = result.where(counts >= self._min_periods)
-
-        if self.center:
-            result = self._center_result(result)
-
-        return result
+        # restore dim order
+        return result.transpose(*self.obj.dims)
 
     @classmethod
     def _reduce_method(cls, func):
@@ -254,19 +257,24 @@ def wrapped_func(self, **kwargs):
 
             axis = self.obj.get_axis_num(self.dim)
 
-            if isinstance(self.obj.data, dask_array_type):
+            padded = self.obj.variable
+            if self.center:
+                shift = (-self.window // 2) + 1
+                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,
                                               window=self.window,
                                               min_count=min_count,
                                               axis=axis)
             else:
-                values = func(self.obj.data, window=self.window,
+                values = func(padded.data, window=self.window,
                               min_count=min_count, axis=axis)
 
-            result = DataArray(values, self.obj.coords)
-
             if self.center:
-                result = self._center_result(result)
+                values = values[valid]
+            result = DataArray(values, self.obj.coords)
 
             return result
         return wrapped_func
@@ -373,6 +381,31 @@ def wrapped_func(self, **kwargs):
             return Dataset(reduced, coords=self.obj.coords)
         return wrapped_func
 
+    def to_dataset(self, window_dim):
+        """
+        Convert this rolling object to xr.Dataset,
+        where the window dimension is stacked as a new dimension
+
+        Parameters
+        ----------
+        window_dim: str
+            New name of the window dimension.
+
+        Returns
+        -------
+        Dataset with variables converted from rolling object.
+        """
+
+        from .dataset import Dataset
+
+        dataset = OrderedDict()
+        for key, da in self.obj.data_vars.items():
+            if self.dim in da.dims:
+                dataset[key] = self.rollings[key].to_dataarray(window_dim)
+            else:
+                dataset[key] = da
+        return Dataset(dataset, coords=self.obj.coords)
+
 
 inject_bottleneck_rolling_methods(DataArrayRolling)
 inject_datasetrolling_methods(DatasetRolling)