[PYTHON] [numpy] Create a moving window matrix from multidimensional time series data

Former story

[\ numpy ] Create a moving window matrix from time series data --Qiita I tried to extend this to multidimensional data.

Thing you want to do

Given an ndarray x, considers the last dimension as a time and creates an ndarray y with the data contained in a moving window with a width of $ w $.

When x is one-dimensional

x = 
\begin{pmatrix}
x_0 & x_1 & \cdots & x_{n-1}
\end{pmatrix}

Then

y = 
\begin{pmatrix}
x_0 & x_1 & \cdots & x_{n-w} \\
x_1 & x_2 & \cdots & x_{n-w+1} \\
\vdots & \vdots & & \vdots \\
x_{w-1} & x_w & \cdots & x_{n-1} \\
\end{pmatrix}

It will be. Also, when x is two-dimensional

x = 
\begin{pmatrix}
x_{0,0} & x_{0,1} & \cdots & x_{0,n-1} \\
\vdots & \vdots & & \vdots \\
x_{m-1,0} & x_{m-1,1} & \cdots & x_{m-1,n-1}
\end{pmatrix}

Then, when slicing in the third dimension

y_{:,:,t} = 
\begin{pmatrix}
x_{0,t} & x_{0,t+1} & \cdots & x_{0,t+w-1} \\
\vdots & \vdots & & \vdots \\
x_{m-1,t} & x_{m-1,t+1} & \cdots & x_{m-1,t+w-1}
\end{pmatrix}

Create an ndarray y such that It is the same even if there are more dimensions.

code

import numpy as np

def moving_window_matrix(x, window_size):
    n = x.shape[-1]
    new_shape = tuple(x.shape[0:-1]) + (window_size, n-window_size+1)
    stride = x.strides[-1]
    return np.lib.stride_tricks.as_strided(
        x,
        shape=new_shape,
        strides=x.strides + (stride,)).copy()

x = np.arange(30).reshape((3,10))
print(x)
window_size = 3
y = moving_window_matrix(x, window_size)
print(y[:, :, 0])
print(y[:, :, 1])
print(y[:, :, 2])
print(y[:, :, 7])