Python 阵列轴上滚动或滑动窗口的广义方法

如何有效地在给定阵列的任意轴上创建滑动窗口阵列？例如，如果我有以下数组：

[[ 0  1  2  3  4]
 [ 5  6  7  8  9]
 [10 11 12 13 14]
 [15 16 17 18 19]
 [20 21 22 23 24]
 [25 26 27 28 29]]

窗口大小为4，我希望能够在第一维度上制作一个滑动窗口，如下所示：

[[[ 0  1  2  3  4]
  [ 5  6  7  8  9]
  [10 11 12 13 14]
  [15 16 17 18 19]]

 [[ 5  6  7  8  9]
  [10 11 12 13 14]
  [15 16 17 18 19]
  [20 21 22 23 24]]

 [[10 11 12 13 14]
  [15 16 17 18 19]
  [20 21 22 23 24]
  [25 26 27 28 29]]]

[[[ 0  1  2  3]
  [ 5  6  7  8]
  [10 11 12 13]
  [15 16 17 18]
  [20 21 22 23]
  [25 26 27 28]]

 [[ 1  2  3  4]
  [ 6  7  8  9]
  [11 12 13 14]
  [16 17 18 19]
  [21 22 23 24]
  [26 27 28 29]]]

但也跨越第二维度，如下所示：

[[[ 0  1  2  3  4]
  [ 5  6  7  8  9]
  [10 11 12 13 14]
  [15 16 17 18 19]]

 [[ 5  6  7  8  9]
  [10 11 12 13 14]
  [15 16 17 18 19]
  [20 21 22 23 24]]

 [[10 11 12 13 14]
  [15 16 17 18 19]
  [20 21 22 23 24]
  [25 26 27 28 29]]]

[[[ 0  1  2  3]
  [ 5  6  7  8]
  [10 11 12 13]
  [15 16 17 18]
  [20 21 22 23]
  [25 26 27 28]]

 [[ 1  2  3  4]
  [ 6  7  8  9]
  [11 12 13 14]
  [16 17 18 19]
  [21 22 23 24]
  [26 27 28 29]]]

---

您可以在固定的时间内高效地构建这样的阵列，而无需使用任何额外的内存。生成的数组将是一个有一些限制的视图，但如果需要连续数组，则始终可以创建一个

以下函数解决了一般问题：

import numpy as np

def as_sliding_window(x, window_size, axis=0, window_axis=None,
                      subok=False, writeable=True):
    """
    Make a sliding window across an axis.

    Uses ``numpy.lib.stride_tricks.as_strided``, similar caveats apply.

    Parameters
    ----------
    x : array_like
        Array from where the sliding window is created.
    window_size: int
        Size of the sliding window.
    axis: int
        Dimension across which the sliding window is created.
    window_axis: int
        New dimension for the sliding window. By default, the new
        dimension is inserted before ``axis``.
    subok: bool
        If True, subclasses are preserved
        (see ``numpy.lib.stride_tricks.as_strided``).
    writeable: bool
        If set to False, the returned array will always be readonly.
        Otherwise it will be writable if the original array was. It
        is advisable to set this to False if possible
        (see ``numpy.lib.stride_tricks.as_strided``).

    Returns
    --------
    sliding_window: ndarray
        View of the given array as a sliding window along ``axis``.
    """
    from numpy.lib.stride_tricks import as_strided
    x = np.asarray(x)
    axis %= x.ndim
    if window_axis is None:
        window_axis = axis
    window_axis %= x.ndim + 1
    # Make shape
    shape = list(x.shape)
    n = shape[axis]
    shape[axis] = window_size
    shape.insert(window_axis, max(n - window_size + 1, 0))
    # Make strides
    strides = list(x.strides)
    strides.insert(window_axis, strides[axis])
    # Make sliding window view
    sliding_window = as_strided(x, shape, strides,
                                subok=subok, writeable=writeable)
    return sliding_window

示例：

x = np.arange(30).reshape((6, 5))
window_size = 4
print(x)
# [[ 0  1  2  3  4]
#  [ 5  6  7  8  9]
#  [10 11 12 13 14]
#  [15 16 17 18 19]
#  [20 21 22 23 24]
#  [25 26 27 28 29]]

print(as_sliding_window(x, window_size))
# [[[ 0  1  2  3  4]
#   [ 5  6  7  8  9]
#   [10 11 12 13 14]
#   [15 16 17 18 19]]
#
#  [[ 5  6  7  8  9]
#   [10 11 12 13 14]
#   [15 16 17 18 19]
#   [20 21 22 23 24]]
#
#  [[10 11 12 13 14]
#   [15 16 17 18 19]
#   [20 21 22 23 24]
#   [25 26 27 28 29]]]

print(as_sliding_window(x, window_size, axis=1, window_axis=0))
# [[[ 0  1  2  3]
#   [ 5  6  7  8]
#   [10 11 12 13]
#   [15 16 17 18]
#   [20 21 22 23]
#   [25 26 27 28]]
#
#  [[ 1  2  3  4]
#   [ 6  7  8  9]
#   [11 12 13 14]
#   [16 17 18 19]
#   [21 22 23 24]
#   [26 27 28 29]]]

# You can make sliding windows of sliding windows
print(as_sliding_window(as_sliding_window(x, window_size), window_size, axis=2).shape)
# (3, 4, 2, 4)

# New dimension can be put at the end with window_axis=-1
print(as_sliding_window(x, window_size, axis=0, window_axis=-1).shape)
# (4, 5, 3)