ENH: Implement sliding window (gh-17394)

* implement sliding_window_view #7753
Test cases are shown in the issue page.

* Add Example Cases

* Add step_size and N-dim support

* Add shape and step_size check. Remove warning.

* Remove shape default
Add step_size default's description.

* Give proper parameter name 'step'

* fix a parameter description mistake

* implement test function for sliding_window_view()

* implement test function for sliding_window_view()

* Fix according to @eric-wieser comments

* Change arange to ogrid in Examples

* remove np.squeeze on return line

* Clarify document to avoid parameter confusion.

* add `writable` and more explanation in docs

* resolve a write conflit

* fixes according to @seberg review

* resolve write hazard

* remove outdated docs.

* change referring according to @mattip.
change 'writeable' to 'readonly' as @seberg suggest.
remove 'step' as @eric-wieser request

* fix test minor error

* DOC: Grammar fixes

* STY: Add missing line break required by PEP8

* + Change readonly parameter to writeable.
+ Update writeable description.
+ Fix a few parameter checks.
+ Other minor improvements.

* Move to new api as proposed by @eric-wieser

- Change api to follow suggestion by Eric Wieser in
  https://github.com/numpy/numpy/pull/10771#issuecomment-524715356
- Update docstring
- Add more tests

* Improve documentation

* Add sliding_window_view to documentation index

* Apply suggestions from code review

Co-authored-by: Eric Wieser <wieser.eric@gmail.com>

* Fix window shape check

* Add `sliding_window_view` to __all__

* Add tests for  error cases

* Add array_function dispatching

* Change dispatcher argument defaults to None

* Simplify array function dispatching

* Added "np." prefix to doctests

* Split tests

* Improved docstring

* Add release note

* Fix docstring formatting

* Fix doctest

* Remove namespacing in documentation indexing

* Improve docstring

* Improved docstring

* Simplified docstring

* Improve docstring to make pseudo code stand out

* Improve docstring

* Add simple application example

* Correct release note

* Improve link with as_strides

* Add note about performance

* Tidy up main doc string

* Make language on performance warning stronger

* Bugfix: pass subok and writeable to as_strided

* Add writeable test

* Add subok test

* Change subok test to use custom array subclass instead of unsupported MaskedArray

* Add version added information

Co-authored-by: Eric Wieser <wieser.eric@gmail.com>

Co-authored-by: Fanjin <fjzeng@ucsd.edu>
Co-authored-by: Fanjin Zeng <Fnjn@users.noreply.github.com>
Co-authored-by: Eric Wieser <wieser.eric@gmail.com>
Co-authored-by: fanjin <fjzeng@outlook.com>

Closes gh-7753

1 files changed, 227 insertions, 1 deletions

diff --git a/numpy/lib/stride_tricks.py b/numpy/lib/stride_tricks.py
index d8a8b325e..a6c0da751 100644
--- a/numpy/lib/stride_tricks.py
+++ b/numpy/lib/stride_tricks.py
@@ -6,9 +6,11 @@ NumPy reference guide.
 
 """
 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,228 @@ 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.
+    
+    .. versionadded:: 1.20.0
+
+    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)