Sliding window of M-by-N shape numpy.ndarray

I have a Numpy array of shape (6,2):

[[ 0, 1],
 [10,11],
 [20,21],
 [30,31],
 [40,41],
 [50,51]]

I need a sliding window with step size 1 and window size 3 like this:

[[ 0, 1,10,11,20,21],
 [10,11,20,21,30,31],
 [20,21,30,31,40,41],
 [30,31,40,41,50,51]]

I’m looking for a Numpy solution. If your solution could parametrise the shape of the original array as well as the window size and step size, that’d be great.

---

I found this related answer Using strides for an efficient moving average filter but I don’t see how to specify the stepsize there and how to collapse the window from the 3d to a continuous 2d array. Also this Rolling or sliding window iterator? but that’s in Python and I’m not sure how efficient that is. Also, it supports elements but does not join them together in the end if each element has multiple features.

Answer

In [1]: import numpy as np

In [2]: a = np.array([[00,01], [10,11], [20,21], [30,31], [40,41], [50,51]])

In [3]: w = np.hstack((a[:-2],a[1:-1],a[2:]))

In [4]: w
Out[4]: 
array([[ 0,  1, 10, 11, 20, 21],
       [10, 11, 20, 21, 30, 31],
       [20, 21, 30, 31, 40, 41],
       [30, 31, 40, 41, 50, 51]])

You could write this in as a function as so:

def window_stack(a, stepsize=1, width=3):
    n = a.shape[0]
    return np.hstack( a[i:1+n+i-width:stepsize] for i in range(0,width) )

---

This doesn’t really depend on the shape of the original array, as long as a.ndim = 2. Note that I never use either lengths in the interactive version. The second dimension of the shape is irrelevant; each row can be as long as you want. Thanks to @Jaime’s suggestion, you can do it without checking the shape at all:

def window_stack(a, stepsize=1, width=3):
    return np.hstack( a[i:1+i-width or None:stepsize] for i in range(0,width) )