Repeat Numpy array by a sliding window

Question

From the following array of shape (6, 3) :

>>> arr
[
    [1, 0, 1],
    [0, 0, 2],
    [1, 2, 0],
    [0, 1, 3],
    [2, 2, 1],
    [2, 0, 2]
]

I'd like to repeat the values according to a sliding window of n=4 , giving a new array of shape (6-n-1, n, 3) :

>>> new_arr
[
    [
        [1, 0, 1],
        [0, 0, 2],
        [1, 2, 0],
        [0, 1, 3]
    ],
    [
        [0, 0, 2],
        [1, 2, 0],
        [0, 1, 3],
        [2, 2, 1]
    ],
    [
        [1, 2, 0],
        [0, 1, 3],
        [2, 2, 1],
        [2, 0, 2]
    ]
]

It is relatively straightforward using a loop, but it gets extremely slow with several million values (instead of 6 in this example) in the initial array.

Is there a faster way to get to new_arr using Numpy primitives?

Answer 1

You can use NumPy, specifically this function (only NumPy >= 1.20.0):

from numpy.lib.stride_tricks import sliding_window_view

new_arr = sliding_window_view(arr, (n, arr.shape[1])).squeeze()