ENH: Implement sliding window

doc/release/upcoming_changes/17394.new_function.rst

@@ -0,0 +1,5 @@
+``sliding_window_view`` provides a sliding window view for numpy arrays
+-----------------------------------------------------------------------
+`numpy.sliding_window_view` constructs views on numpy arrays that offer a
+sliding or moving window access to the array. This allows for the simple
+implementation of certain algorithms, such as running means.

doc/source/reference/routines.indexing.rst

@@ -42,6 +42,7 @@ Indexing-like operations
    diag
    diagonal
    select
+   sliding_window_view
    lib.stride_tricks.as_strided
 
 Inserting data into arrays

numpy/lib/stride_tricks.py

@@ -6,9 +6,11 @@
 
 """
 import numpy as np
+from numpy.core.numeric import normalize_axis_tuple
 from numpy.core.overrides import array_function_dispatch, set_module
 
-__all__ = ['broadcast_to', 'broadcast_arrays', 'broadcast_shapes']
+__all__ = ['broadcast_to', 'broadcast_arrays', 'broadcast_shapes',
+           'sliding_window_view']
 
 
 class DummyArray:
@@ -67,6 +69,8 @@ def as_strided(x, shape=None, strides=None, subok=False, writeable=True):
     --------
     broadcast_to: broadcast an array to a given shape.
     reshape : reshape an array.
+    sliding_window_view: userfriendly and safe function for the creation of
+        sliding window views.
 
     Notes
     -----
@@ -111,6 +115,226 @@ def as_strided(x, shape=None, strides=None, subok=False, writeable=True):
     return view
 
 
+def _sliding_window_view_dispatcher(x, window_shape, axis=None, *,
+                                    subok=None, writeable=None):
+    return (x,)
+
+
+@array_function_dispatch(_sliding_window_view_dispatcher, module='numpy')
+def sliding_window_view(x, window_shape, axis=None, *,
+                        subok=False, writeable=False):
+    """
+    Create a sliding window view into the array with the given window shape.
+
+    Also known as rolling or moving window, the window slides across all
+    dimensions of the array and extracts subsets of the array at all window
+    positions.
+
+    Parameters
+    ----------
+    x : array_like
+        Array to create the sliding window view from.
+    window_shape : int or tuple of int
+        Size of window over each axis that takes part in the sliding window.
+        If `axis` is not present, must have same length as the number of input
+        array dimensions. Single integers `i` are treated as if they were the
+        tuple `(i,)`.
+    axis : int or tuple of int, optional
+        Axis or axes along which the sliding window is applied.
+        By default, the sliding window is applied to all axes and
+        `window_shape[i]` will refer to axis `i` of `x`.
+        If `axis` is given as a `tuple of int`, `window_shape[i]` will refer to
+        the axis `axis[i]` of `x`.
+        Single integers `i` are treated as if they were the tuple `(i,)`.
+    subok : bool, optional
+        If True, sub-classes will be passed-through, otherwise the returned
+        array will be forced to be a base-class array (default).
+    writeable : bool, optional
+        When true, allow writing to the returned view. The default is false,
+        as this should be used with caution: the returned view contains the
+        same memory location multiple times, so writing to one location will
+        cause others to change.
+
+    Returns
+    -------
+    view : ndarray
+        Sliding window view of the array. The sliding window dimensions are
+        inserted at the end, and the original dimensions are trimmed as
+        required by the size of the sliding window.
+        That is, ``view.shape = x_shape_trimmed + window_shape``, where
+        ``x_shape_trimmed`` is ``x.shape`` with every entry reduced by one less
+        than the corresponding window size.
+
+    See Also
+    --------
+    lib.stride_tricks.as_strided: A lower-level and less safe routine for
+        creating arbitrary views from custom shape and strides.
+    broadcast_to: broadcast an array to a given shape.
+
+    Notes
+    -----
+    For many applications using a sliding window view can be convenient, but
+    potentially very slow. Often specialized solutions exist, for example:
+
+    - `scipy.signal.fftconvolve`
+
+    - filtering functions in `scipy.ndimage`
+
+    - moving window functions provided by
+      `bottleneck <https://github.com/pydata/bottleneck>`_.
+
+    As a rough estimate, a sliding window approach with an input size of `N`
+    and a window size of `W` will scale as `O(N*W)` where frequently a special
+    algorithm can achieve `O(N)`. That means that the sliding window variant
+    for a window size of 100 can be a 100 times slower than a more specialized
+    version.
+
+    Nevertheless, for small window sizes, when no custom algorithm exists, or
+    as a prototyping and developing tool, this function can be a good solution.
+
+    Examples
+    --------
+    >>> x = np.arange(6)
+    >>> x.shape
+    (6,)
+    >>> v = np.sliding_window_view(x, 3)
+    >>> v.shape
+    (4, 3)
+    >>> v
+    array([[0, 1, 2],
+           [1, 2, 3],
+           [2, 3, 4],
+           [3, 4, 5]])
+
+    This also works in more dimensions, e.g.
+
+    >>> i, j = np.ogrid[:3, :4]
+    >>> x = 10*i + j
+    >>> x.shape
+    (3, 4)
+    >>> x
+    array([[ 0,  1,  2,  3],
+           [10, 11, 12, 13],
+           [20, 21, 22, 23]])
+    >>> shape = (2,2)
+    >>> v = np.sliding_window_view(x, shape)
+    >>> v.shape
+    (2, 3, 2, 2)
+    >>> v
+    array([[[[ 0,  1],
+             [10, 11]],
+            [[ 1,  2],
+             [11, 12]],
+            [[ 2,  3],
+             [12, 13]]],
+           [[[10, 11],
+             [20, 21]],
+            [[11, 12],
+             [21, 22]],
+            [[12, 13],
+             [22, 23]]]])
+
+    The axis can be specified explicitly:
+
+    >>> v = np.sliding_window_view(x, 3, 0)
+    >>> v.shape
+    (1, 4, 3)
+    >>> v
+    array([[[ 0, 10, 20],
+            [ 1, 11, 21],
+            [ 2, 12, 22],
+            [ 3, 13, 23]]])
+
+    The same axis can be used several times. In that case, every use reduces
+    the corresponding original dimension:
+
+    >>> v = np.sliding_window_view(x, (2, 3), (1, 1))
+    >>> v.shape
+    (3, 1, 2, 3)
+    >>> v
+    array([[[[ 0,  1,  2],
+             [ 1,  2,  3]]],
+           [[[10, 11, 12],
+             [11, 12, 13]]],
+           [[[20, 21, 22],
+             [21, 22, 23]]]])
+
+    Combining with stepped slicing (`::step`), this can be used to take sliding
+    views which skip elements:
+
+    >>> x = np.arange(7)
+    >>> np.sliding_window_view(x, 5)[:, ::2]
+    array([[0, 2, 4],
+           [1, 3, 5],
+           [2, 4, 6]])
+
+    or views which move by multiple elements
+
+    >>> x = np.arange(7)
+    >>> np.sliding_window_view(x, 3)[::2, :]
+    array([[0, 1, 2],
+           [2, 3, 4],
+           [4, 5, 6]])
+
+    A common application of `sliding_window_view` is the calculation of running
+    statistics. The simplest example is the
+    `moving average <https://en.wikipedia.org/wiki/Moving_average>`_:
+
+    >>> x = np.arange(6)
+    >>> x.shape
+    (6,)
+    >>> v = np.sliding_window_view(x, 3)
+    >>> v.shape
+    (4, 3)
+    >>> v
+    array([[0, 1, 2],
+           [1, 2, 3],
+           [2, 3, 4],
+           [3, 4, 5]])
+    >>> moving_average = v.mean(axis=-1)
+    >>> moving_average
+    array([1., 2., 3., 4.])
+
+    Note that a sliding window approach is often **not** optimal (see Notes).
+    """
+    window_shape = (tuple(window_shape)
+                    if np.iterable(window_shape)
+                    else (window_shape,))
+    # first convert input to array, possibly keeping subclass
+    x = np.array(x, copy=False, subok=subok)
+
+    window_shape_array = np.array(window_shape)
+    if np.any(window_shape_array < 0):
+        raise ValueError('`window_shape` cannot contain negative values')
+
+    if axis is None:
+        axis = tuple(range(x.ndim))
+        if len(window_shape) != len(axis):
+            raise ValueError(f'Since axis is `None`, must provide '
+                             f'window_shape for all dimensions of `x`; '
+                             f'got {len(window_shape)} window_shape elements '
+                             f'and `x.ndim` is {x.ndim}.')
+    else:
+        axis = normalize_axis_tuple(axis, x.ndim, allow_duplicate=True)
+        if len(window_shape) != len(axis):
+            raise ValueError(f'Must provide matching length window_shape and '
+                             f'axis; got {len(window_shape)} window_shape '
+                             f'elements and {len(axis)} axes elements.')
+
+    out_strides = x.strides + tuple(x.strides[ax] for ax in axis)
+
+    # note: same axis can be windowed repeatedly
+    x_shape_trimmed = list(x.shape)
+    for ax, dim in zip(axis, window_shape):
+        if x_shape_trimmed[ax] < dim:
+            raise ValueError(
+                'window shape cannot be larger than input array shape')
+        x_shape_trimmed[ax] -= dim - 1
+    out_shape = tuple(x_shape_trimmed) + window_shape
+    return as_strided(x, strides=out_strides, shape=out_shape,
+                      subok=subok, writeable=writeable)
+
+
 def _broadcast_to(array, shape, subok, readonly):
     shape = tuple(shape) if np.iterable(shape) else (shape,)
     array = np.array(array, copy=False, subok=subok)

numpy/lib/tests/test_stride_tricks.py

@@ -6,8 +6,10 @@
     )
 from numpy.lib.stride_tricks import (
     as_strided, broadcast_arrays, _broadcast_shape, broadcast_to,
-    broadcast_shapes,
+    broadcast_shapes, sliding_window_view,
     )
+import pytest
+
 
 def assert_shapes_correct(input_shapes, expected_shape):
     # Broadcast a list of arrays with the given input shapes and check the
@@ -394,6 +396,109 @@ def test_as_strided():
     assert_equal(a.dtype, a_view.dtype)
     assert_array_equal([r] * 3, a_view)
 
+
+class TestSlidingWindowView:
+    def test_1d(self):
+        arr = np.arange(5)
+        arr_view = sliding_window_view(arr, 2)
+        expected = np.array([[0, 1],
+                             [1, 2],
+                             [2, 3],
+                             [3, 4]])
+        assert_array_equal(arr_view, expected)
+
+    def test_2d(self):
+        i, j = np.ogrid[:3, :4]
+        arr = 10*i + j
+        shape = (2, 2)
+        arr_view = sliding_window_view(arr, shape)
+        expected = np.array([[[[0, 1], [10, 11]],
+                              [[1, 2], [11, 12]],
+                              [[2, 3], [12, 13]]],
+                             [[[10, 11], [20, 21]],
+                              [[11, 12], [21, 22]],
+                              [[12, 13], [22, 23]]]])
+        assert_array_equal(arr_view, expected)
+
+    def test_2d_with_axis(self):
+        i, j = np.ogrid[:3, :4]
+        arr = 10*i + j
+        arr_view = sliding_window_view(arr, 3, 0)
+        expected = np.array([[[0, 10, 20],
+                              [1, 11, 21],
+                              [2, 12, 22],
+                              [3, 13, 23]]])
+        assert_array_equal(arr_view, expected)
+
+    def test_2d_repeated_axis(self):
+        i, j = np.ogrid[:3, :4]
+        arr = 10*i + j
+        arr_view = sliding_window_view(arr, (2, 3), (1, 1))
+        expected = np.array([[[[0, 1, 2],
+                               [1, 2, 3]]],
+                             [[[10, 11, 12],
+                               [11, 12, 13]]],
+                             [[[20, 21, 22],
+                               [21, 22, 23]]]])
+        assert_array_equal(arr_view, expected)
+
+    def test_2d_without_axis(self):
+        i, j = np.ogrid[:4, :4]
+        arr = 10*i + j
+        shape = (2, 3)
+        arr_view = sliding_window_view(arr, shape)
+        expected = np.array([[[[0, 1, 2], [10, 11, 12]],
+                              [[1, 2, 3], [11, 12, 13]]],
+                             [[[10, 11, 12], [20, 21, 22]],
+                              [[11, 12, 13], [21, 22, 23]]],
+                             [[[20, 21, 22], [30, 31, 32]],
+                              [[21, 22, 23], [31, 32, 33]]]])
+        assert_array_equal(arr_view, expected)
+
+    def test_errors(self):
+        i, j = np.ogrid[:4, :4]
+        arr = 10*i + j
+        with pytest.raises(ValueError, match='cannot contain negative values'):
+            sliding_window_view(arr, (-1, 3))
+        with pytest.raises(
+                ValueError,
+                match='must provide window_shape for all dimensions of `x`'):
+            sliding_window_view(arr, (1,))
+        with pytest.raises(
+                ValueError,
+                match='Must provide matching length window_shape and axis'):
+            sliding_window_view(arr, (1, 3, 4), axis=(0, 1))
+        with pytest.raises(
+                ValueError,
+                match='window shape cannot be larger than input array'):
+            sliding_window_view(arr, (5, 5))
+
+    def test_writeable(self):
+        arr = np.arange(5)
+        view = sliding_window_view(arr, 2, writeable=False)
+        assert_(not view.flags.writeable)
+        with pytest.raises(
+                ValueError,
+                match='assignment destination is read-only'):
+            view[0, 0] = 3
+        view = sliding_window_view(arr, 2, writeable=True)
+        assert_(view.flags.writeable)
+        view[0, 1] = 3
+        assert_array_equal(arr, np.array([0, 3, 2, 3, 4]))
+
+    def test_subok(self):
+        class MyArray(np.ndarray):
+            pass
+
+        arr = np.arange(5).view(MyArray)
+        assert_(not isinstance(sliding_window_view(arr, 2,
+                                                   subok=False),
+                               MyArray))
+        assert_(isinstance(sliding_window_view(arr, 2, subok=True), MyArray))
+        # Default behavior
+        assert_(not isinstance(sliding_window_view(arr, 2), MyArray))
+
+
 def as_strided_writeable():
     arr = np.ones(10)
     view = as_strided(arr, writeable=False)
@@ -496,7 +601,7 @@ def test_writeable():
             # check: no warning emitted
             assert_equal(result.flags.writeable, True)
             result[:] = 0
-            
+
     # keep readonly input readonly
     original.flags.writeable = False
     _, result = broadcast_arrays(0, original)