Sliding window function for NumPy arrays

sliding_window.py

# from
# https://gist.githubusercontent.com/nils-werner/9d321441006b112a4b116a8387c2280c/raw/489137f314c022ca08f0c5089319ded2eb8c5ae1/sliding_window.py
def sliding_window(size, data, step=1, padded=False, axis=-1, copy=True):
    """Calculate a sliding window over an array.

    Note:
        If the array cannot be windowed evenly, it is truncated at the end.

    Args:
        data (np.ndarray): The array to be windowed.
        size (int): The window size.
        step (int): The number of elements between windows. Defaults to 1.
        axis (int): The axis to slide over. Defaults to the last axis.
        copy (bool): Return a copy instead of a view.

    Returns:
        np.ndarray: A matrix where the last dimension corresponds to the
        window.

    Examples:
        >>> a = np.array([1, 2, 3, 4, 5])
        >>> sliding_window(3, a)
        array([[1, 2, 3],
               [2, 3, 4],
               [3, 4, 5]])
        >>> sliding_window(3, a, step=2)
        array([[1, 2, 3],
               [3, 4, 5]])
    """
    if axis >= data.ndim:
        raise ValueError("axis value out of range")
    if step < 1:
        raise ValueError("step size must be positive")
    if size > data.shape[axis]:
        raise ValueError("window size may not exceed size of selected axis")

    shape = list(data.shape)
    shape[axis] = np.floor(data.shape[axis] / step - size / step + 1).astype(int)
    shape.append(size)

    strides = list(data.strides)
    strides[axis] *= step
    strides.append(data.strides[axis])

    strided = np.lib.stride_tricks.as_strided(data, shape=shape, strides=strides)

    if copy:
        return strided.copy()
    else:
        return strided