Repository: dask/dask
Pull request: Add sliding_window_view
- [x] Closes #4659
- [x] Tests added / passed
- [x] needs docstring
- [x] match numpy's default kwargs
- [x] copy over error checking from numpy
- [x] Passes `black dask` / `flake8 dask`


Problem statement:
Support sliding window computations
Consider the following code:

```python
import dask.array as da
import numpy as np

d = da.arange(8, chunks=4)

g = d.map_overlap(np.mean, depth=0, boundary=0, trim=False, keepdims=True)
```

After computing `g`, this will return the mean of the two chunks `[0:4]` and `[4:8]`. Currently, there seems to be no way to specify a step size different (smaller) than the chunk size for each dimension. One could want a step size of 2, result in the mean for ranges `[0:4]`, `[2:6]`, `[4:8]`.

Being able to specify a step size is very useful for filtering images, for example, as the user may be fine with computing every the filter for only every second or third pixel to reduce compute time. The user may want an arbitrary step size, which I think will normally be anything from `1..len(chunk[dim])`.

I would be interesting in knowing what options do we currently have (if any), otherwise, what people think may be a good direction for the implementation of a step size.

cc @mrocklin @jakirkham

Hints from issue discussion:
I've renamed this issue to "Support sliding window comptutations" (I hope that that's ok).  The map_overlap function is more specifically for mapping a function over chunks of data, my guess is that it would be used internally by some sort of sliding window function, but may not be the user level API directly.

Today, NumPy lacks a sliding window API.  This is discussed in more depth here: https://github.com/numpy/numpy/issues/7753

If such an API existed, then we might consider copying it.  Given that it doesn't exist, we might consider creating our own, and collaborating with upstream NumPy to make sure that they're ok with it in case they ever go in that direction.

If I were to do this today I would probably make a block-wise function that did a sliding window computation, probably using either NumPy stride tricks (a hack described in that issue and known to expert numpy users) or just for loops with Numba, and then I would apply that function with `map_overlap`
Thanks for your input @mrocklin. I'm definitely gonna take a look at the NumPy discussion in detail.

I have a feeling you're right, we're probably gonna need the block-wise function, or Numba as a last resort. I'm particularly concerned about the memory footprint, certainly making a copy of every window+boundaries to use it as a chunk will be too expensive.
I think that we can use either Numba or stride tricks to do an efficient sliding window computation in low memory.  

I think that we can use map_overlap to handle the boundaries between large blocks.

I think that with both of these we can achieve what you want in small space.
I'll continue to investigate this, thanks for the suggestions @mrocklin.
There's a sliding window PR https://github.com/numpy/numpy/pull/10771, it seems to be what we want. The progress has stalled for quite some time, the author mentioned he would pick it up again, but that hasn't happened yet, presumable he's got no time for that. I can probably suggest there that I continue with it. Do you think that looks good @mrocklin ?

I'd nevertheless have to find a more immediate solution, since we'll need the same for CuPy, and I don't think we can get that PR in NumPy, do the work for CuPy as well and have it merged in less than 1-2 months.
I think it would be good for that PR to be finished.  In principle having anyone finish it would be good.

Putting on my NVIDIA hat I don't currently know where this task fits into our current priorities.  It's not clearly tied to any particular task that I know of, so I would prioritize is somewhat low.  That being said, if you're excited about doing it then please don't let me get in the way :)

Commit message: add breaking test

Patch to review:
diff --git a/dask/array/__init__.py b/dask/array/__init__.py
index 9d9d5ef7121..9d061f5b3f4 100644
--- a/dask/array/__init__.py
+++ b/dask/array/__init__.py
@@ -219,7 +219,7 @@
     from .percentile import percentile
     from . import ma
     from . import random, linalg, overlap, fft, backends
-    from .overlap import map_overlap
+    from .overlap import map_overlap, sliding_window_view
     from .wrap import ones, zeros, empty, full
     from .creation import ones_like, zeros_like, empty_like, full_like
     from .rechunk import rechunk
diff --git a/dask/array/numpy_compat.py b/dask/array/numpy_compat.py
index d3f4400abf3..cbaf709bdce 100644
--- a/dask/array/numpy_compat.py
+++ b/dask/array/numpy_compat.py
@@ -94,7 +94,7 @@ def map_reduce(
         f_map=lambda x, **kwargs: x,
         f_reduce=lambda x, **kwargs: x,
         f_kwargs=lambda **kwargs: kwargs,
-        **kwargs
+        **kwargs,
     ):
         """
         Iterate over the nested list, applying:
@@ -206,3 +206,95 @@ def rollaxis(a, axis, start=0):
     axes.remove(axis)
     axes.insert(start, axis)
     return a.transpose(axes)
+
+
+if _numpy_120:
+    sliding_window_view = np.lib.stride_tricks.sliding_window_view
+else:
+
+    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.
+        """
+        from numpy.core.numeric import normalize_axis_tuple
+
+        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 np.lib.stride_tricks.as_strided(
+            x, strides=out_strides, shape=out_shape, subok=subok, writeable=writeable
+        )
diff --git a/dask/array/overlap.py b/dask/array/overlap.py
index 3bd61cbf220..45b4d3d059e 100644
--- a/dask/array/overlap.py
+++ b/dask/array/overlap.py
@@ -469,6 +469,16 @@ def ensure_minimum_chunksize(size, chunks):
     return tuple(output)
 
 
+def rechunk_for_overlap(x, depth):
+    depth2 = coerce_depth(x.ndim, depth)
+    # rechunk if new chunks are needed to fit depth in every chunk
+    depths = [max(d) if isinstance(d, tuple) else d for d in depth2.values()]
+    new_chunks = tuple(
+        ensure_minimum_chunksize(size, c) for size, c in zip(depths, x.chunks)
+    )
+    return x.rechunk(new_chunks)
+
+
 def overlap(x, depth, boundary):
     """Share boundaries between neighboring blocks
 
@@ -527,13 +537,7 @@ def overlap(x, depth, boundary):
     depth2 = coerce_depth(x.ndim, depth)
     boundary2 = coerce_boundary(x.ndim, boundary)
 
-    # rechunk if new chunks are needed to fit depth in every chunk
-    depths = [max(d) if isinstance(d, tuple) else d for d in depth2.values()]
-    new_chunks = tuple(
-        ensure_minimum_chunksize(size, c) for size, c in zip(depths, x.chunks)
-    )
-    x1 = x.rechunk(new_chunks)
-
+    x1 = rechunk_for_overlap(x, depth2)
     x2 = boundaries(x1, depth2, boundary2)
     x3 = overlap_internal(x2, depth2)
     trim = dict(
@@ -841,3 +845,63 @@ def coerce_boundary(ndim, boundary):
     if isinstance(boundary, dict):
         boundary = {ax: boundary.get(ax, default) for ax in range(ndim)}
     return boundary
+
+
+def _window_view(arr, window_shape, axis, fix_first_block, block_info):
+    """ core sliding window function"""
+    from .numpy_compat import sliding_window_view
+
+    block_info = block_info[0]
+    slicer = [slice(None, None)] * arr.ndim
+    for ax, window, fix_first in zip(axis, window_shape, fix_first_block):
+        start, stop = None, None
+        is_last_chunk = (
+            block_info["num-chunks"][ax] - 1 - block_info["chunk-location"][ax]
+        ) == 0
+
+        locs = block_info["array-location"][ax]
+        size = locs[1] - locs[0]
+        if not is_last_chunk:
+            stop = -1 * max(window // 2 - 1, 1)
+
+        # when depth == chunk size of first chunk along some axis
+        # we need to make some confusing empirical adjustments.
+        # could be avoided if we rechunked so that min
+        # chunk size is depth+1 (which would also be more efficient)
+        # 1) don't trim the first block along that axis
+        if size == window and block_info["chunk-location"][ax] == 0:
+            stop = None
+        # 2) we need to skip the first "row" of the second block
+        if fix_first and block_info["chunk-location"][ax] == 1:
+            start = 1
+        slicer[ax] = slice(start, stop)
+
+    view = sliding_window_view(arr[tuple(slicer)], window_shape, axis)
+    return view
+
+
+def sliding_window_view(arr, window_shape, axis):
+    """ Dask wrapper for numpy.lib.stride_tricks.sliding_window_view."""
+
+    depths = {ax: (window, 0) for ax, window in zip(axis, window_shape)}
+
+    # The only way (that I've found) to set chunks in the map_overlap call
+    # is to do the rechunking that overlap would do in any case.
+    arr = rechunk_for_overlap(arr, depths)
+
+    return map_overlap(
+        _window_view,
+        arr,
+        depth=depths,
+        boundary="none",
+        meta=arr._meta,
+        new_axis=range(arr.ndim, arr.ndim + len(axis)),
+        chunks=arr.chunks + tuple((window,) for window in window_shape),
+        trim=False,
+        align_arrays=False,
+        window_shape=window_shape,
+        axis=axis,
+        fix_first_block=tuple(
+            arr.chunks[ax][0] == window for ax, window in zip(axis, window_shape)
+        ),
+    )
diff --git a/dask/array/tests/test_overlap.py b/dask/array/tests/test_overlap.py
index b8fa4885157..f556683153a 100644
--- a/dask/array/tests/test_overlap.py
+++ b/dask/array/tests/test_overlap.py
@@ -784,3 +784,15 @@ def test_ensure_minimum_chunksize_raises_error():
     chunks = (5, 2, 1, 1)
     with pytest.raises(ValueError, match="overlapping depth 10 is larger than"):
         ensure_minimum_chunksize(10, chunks)
+
+
+@pytest.mark.parametrize("window_shape, axis", [((3, 2), (1, 2))])
+def test_sliding_window_view(window_shape, axis):
+    from ..numpy_compat import sliding_window_view as np_sliding_window_view
+    from ..overlap import sliding_window_view
+
+    arr = da.from_array(np.arange(60.0).reshape(3, 4, 5), chunks=(2, 3, 4))
+    actual = sliding_window_view(arr, window_shape, axis)
+    expected = np_sliding_window_view(arr.compute(), window_shape, axis)
+    assert expected.shape == actual.shape
+    assert_array_equal(expected, actual)
diff --git a/docs/source/array-api.rst b/docs/source/array-api.rst
index b516aae166c..e3cb9caf800 100644
--- a/docs/source/array-api.rst
+++ b/docs/source/array-api.rst
@@ -355,6 +355,7 @@ Slightly Overlapping Computations
 .. autosummary::
    overlap.overlap
    overlap.map_overlap
+   overlap.sliding_window_view
    overlap.trim_internal
    overlap.trim_overlap
 


Code review comments:
diff_hunk: @@ -841,3 +845,63 @@ def coerce_boundary(ndim, boundary):
     if isinstance(boundary, dict):
         boundary = {ax: boundary.get(ax, default) for ax in range(ndim)}
     return boundary
+
+
+def _window_view(arr, window_shape, axis, fix_first_block, block_info):
+    """ core sliding window function"""
+    from .numpy_compat import sliding_window_view
+
+    block_info = block_info[0]
+    slicer = [slice(None, None)] * arr.ndim
+    for ax, window, fix_first in zip(axis, window_shape, fix_first_block):
+        start, stop = None, None
+        is_last_chunk = (
+            block_info["num-chunks"][ax] - 1 - block_info["chunk-location"][ax]
+        ) == 0
+
+        locs = block_info["array-location"][ax]
+        size = locs[1] - locs[0]
+        if not is_last_chunk:
+            stop = -1 * max(window // 2 - 1, 1)
+
+        # when depth == chunk size of first chunk along some axis
+        # we need to make some confusing empirical adjustments.
+        # could be avoided if we rechunked so that min
+        # chunk size is depth+1 (which would also be more efficient)
+        # 1) don't trim the first block along that axis
+        if size == window and block_info["chunk-location"][ax] == 0:
+            stop = None
+        # 2) we need to skip the first "row" of the second block
+        if fix_first and block_info["chunk-location"][ax] == 1:
+            start = 1
+        slicer[ax] = slice(start, stop)
+
+    view = sliding_window_view(arr[tuple(slicer)], window_shape, axis)
+    return view
+
+
+def sliding_window_view(arr, window_shape, axis):
+    """ Dask wrapper for numpy.lib.stride_tricks.sliding_window_view."""
+
+    depths = {ax: (window, 0) for ax, window in zip(axis, window_shape)}
+
+    # The only way (that I've found) to set chunks in the map_overlap call
+    # is to do the rechunking that overlap would do in any case.
+    arr = rechunk_for_overlap(arr, depths)
+
+    return map_overlap(
+        _window_view,
+        arr,
+        depth=depths,
+        boundary="none",
+        meta=arr._meta,
+        new_axis=range(arr.ndim, arr.ndim + len(axis)),
+        chunks=arr.chunks + tuple((window,) for window in window_shape),
path: dask/array/overlap.py
text: @author:
This is wrong. Needs to account for
```
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.
```

right now code assumes view.shape = x.shape + window_shape . Final computed result is correct though :)