| 839 | |
| 840 | @derived_from(np.lib.stride_tricks) |
| 841 | def sliding_window_view(x, window_shape, axis=None, automatic_rechunk=True): |
| 842 | window_shape = tuple(window_shape) if np.iterable(window_shape) else (window_shape,) |
| 843 | |
| 844 | window_shape_array = np.array(window_shape) |
| 845 | if np.any(window_shape_array <= 0): |
| 846 | raise ValueError("`window_shape` must contain values > 0") |
| 847 | |
| 848 | if axis is None: |
| 849 | axis = tuple(range(x.ndim)) |
| 850 | if len(window_shape) != len(axis): |
| 851 | raise ValueError( |
| 852 | f"Since axis is `None`, must provide " |
| 853 | f"window_shape for all dimensions of `x`; " |
| 854 | f"got {len(window_shape)} window_shape elements " |
| 855 | f"and `x.ndim` is {x.ndim}." |
| 856 | ) |
| 857 | else: |
| 858 | axis = normalize_axis_tuple(axis, x.ndim, allow_duplicate=True) |
| 859 | if len(window_shape) != len(axis): |
| 860 | raise ValueError( |
| 861 | f"Must provide matching length window_shape and " |
| 862 | f"axis; got {len(window_shape)} window_shape " |
| 863 | f"elements and {len(axis)} axes elements." |
| 864 | ) |
| 865 | |
| 866 | depths = [0] * x.ndim |
| 867 | for ax, window in zip(axis, window_shape): |
| 868 | depths[ax] += window - 1 |
| 869 | |
| 870 | # Ensure that each chunk is big enough to leave at least a size-1 chunk |
| 871 | # after windowing (this is only really necessary for the last chunk). |
| 872 | safe_chunks = list( |
| 873 | ensure_minimum_chunksize(d + 1, c) for d, c in zip(depths, x.chunks) |
| 874 | ) |
| 875 | if automatic_rechunk: |
| 876 | safe_chunks = [ |
| 877 | s if d != 0 else c for d, c, s in zip(depths, x.chunks, safe_chunks) |
| 878 | ] |
| 879 | # safe chunks is our output chunks, so add the new dimensions |
| 880 | safe_chunks.extend([(w,) for w in window_shape]) |
| 881 | max_chunk = reduce(mul, map(max, x.chunks)) |
| 882 | new_chunks = _calculate_new_chunksizes( |
| 883 | x.chunks, |
| 884 | safe_chunks.copy(), |
| 885 | {i for i, d in enumerate(depths) if d == 0}, |
| 886 | max_chunk, |
| 887 | ) |
| 888 | x = x.rechunk(tuple(new_chunks)) |
| 889 | else: |
| 890 | x = x.rechunk(tuple(safe_chunks)) |
| 891 | |
| 892 | # result.shape = x_shape_trimmed + window_shape, |
| 893 | # where x_shape_trimmed is x.shape with every entry |
| 894 | # reduced by one less than the corresponding window size. |
| 895 | # trim chunks to match x_shape_trimmed |
| 896 | newchunks = tuple(c[:-1] + (c[-1] - d,) for d, c in zip(depths, x.chunks)) + tuple( |
| 897 | (window,) for window in window_shape |
| 898 | ) |