[test-upstream] Actually use dask_array_compat.

Author: dcherian

xarray/core/dask_array_compat.py

@@ -152,67 +152,73 @@ def ensure_minimum_chunksize(size, chunks):
         return tuple(output)
 
 
-def sliding_window_view(x, window_shape, axis=None):
-    from dask.array.overlap import map_overlap
-    from numpy.core.numeric import normalize_axis_tuple  # type: ignore
+if LooseVersion(dask_version) > LooseVersion("2021.03.0"):
+    sliding_window_view = da.lib.stride_tricks.sliding_window_view
+else:
 
-    from .npcompat import sliding_window_view as _np_sliding_window_view
+    def sliding_window_view(x, window_shape, axis=None):
+        from dask.array.overlap import map_overlap
+        from numpy.core.numeric import normalize_axis_tuple  # type: ignore
 
-    window_shape = tuple(window_shape) if np.iterable(window_shape) else (window_shape,)
+        from .npcompat import sliding_window_view as _np_sliding_window_view
 
-    window_shape_array = np.array(window_shape)
-    if np.any(window_shape_array <= 0):
-        raise ValueError("`window_shape` must contain positive values")
+        window_shape = (
+            tuple(window_shape) if np.iterable(window_shape) else (window_shape,)
+        )
 
-    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."
-            )
+        window_shape_array = np.array(window_shape)
+        if np.any(window_shape_array <= 0):
+            raise ValueError("`window_shape` must contain positive values")
 
-    depths = [0] * x.ndim
-    for ax, window in zip(axis, window_shape):
-        depths[ax] += window - 1
-
-    # Ensure that each chunk is big enough to leave at least a size-1 chunk
-    # after windowing (this is only really necessary for the last chunk).
-    safe_chunks = tuple(
-        ensure_minimum_chunksize(d + 1, c) for d, c in zip(depths, x.chunks)
-    )
-    x = x.rechunk(safe_chunks)
-
-    # result.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.
-    # trim chunks to match x_shape_trimmed
-    newchunks = tuple(c[:-1] + (c[-1] - d,) for d, c in zip(depths, x.chunks)) + tuple(
-        (window,) for window in window_shape
-    )
-
-    kwargs = dict(
-        depth=tuple((0, d) for d in depths),  # Overlap on +ve side only
-        boundary="none",
-        meta=x._meta,
-        new_axis=range(x.ndim, x.ndim + len(axis)),
-        chunks=newchunks,
-        trim=False,
-        window_shape=window_shape,
-        axis=axis,
-    )
-    # map_overlap's signature changed in https://github.com/dask/dask/pull/6165
-    if LooseVersion(dask_version) > "2.18.0":
-        return map_overlap(_np_sliding_window_view, x, align_arrays=False, **kwargs)
-    else:
-        return map_overlap(x, _np_sliding_window_view, **kwargs)
+        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."
+                )
+
+        depths = [0] * x.ndim
+        for ax, window in zip(axis, window_shape):
+            depths[ax] += window - 1
+
+        # Ensure that each chunk is big enough to leave at least a size-1 chunk
+        # after windowing (this is only really necessary for the last chunk).
+        safe_chunks = tuple(
+            ensure_minimum_chunksize(d + 1, c) for d, c in zip(depths, x.chunks)
+        )
+        x = x.rechunk(safe_chunks)
+
+        # result.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.
+        # trim chunks to match x_shape_trimmed
+        newchunks = tuple(
+            c[:-1] + (c[-1] - d,) for d, c in zip(depths, x.chunks)
+        ) + tuple((window,) for window in window_shape)
+
+        kwargs = dict(
+            depth=tuple((0, d) for d in depths),  # Overlap on +ve side only
+            boundary="none",
+            meta=x._meta,
+            new_axis=range(x.ndim, x.ndim + len(axis)),
+            chunks=newchunks,
+            trim=False,
+            window_shape=window_shape,
+            axis=axis,
+        )
+        # map_overlap's signature changed in https://github.com/dask/dask/pull/6165
+        if LooseVersion(dask_version) > "2.18.0":
+            return map_overlap(_np_sliding_window_view, x, align_arrays=False, **kwargs)
+        else:
+            return map_overlap(x, _np_sliding_window_view, **kwargs)