Rolling windows for ndarrays

Question

I have a ndimensional array with shape (30,2,2) and 2000 elements. So my final array is of shape (2000, 30, 2, 2). I now want to stack rolling 200 elements in a new array. So I assume my final array will look something like (1801, 200, 30, 2, 2) where each element in 1800 has 200 samples of (30,2,2) arrays. How do you create this rolling window in python. I have tried using vstack but not entirely sure how I achieve my desired results.

import numpy as np

input = np.zeros((2000, 30, 2, 2))
desired_output = np.zeros((1801, 200, 30, 2, 2))

Answer 1

We can leverage np.lib.stride_tricks.as_strided based scikit-image's view_as_windows to get sliding windows and hence solve our case here. More info on use of as_strided based view_as_windows.

The output would be a view into the input array and as such, it's virtually free on runtime. The implementation would be -

from skimage.util.shape import view_as_windows

out = view_as_windows(input,(200,1,1,1))[...,0,0,0].transpose(0,4,1,2,3)

Answer 2

You can slice the first dimension and stack window sized arrays together using numpy.stack:

import numpy as np

a = np.zeros((2000, 30, 2, 2))

ws = 200
n = a.shape[0]

r = np.stack([a[i:i+ws] for i in range(n-ws+1)])
r.shape
(1801, 200, 30, 2, 2)