Rolling window over 2D array with specific spacing along one dimension

Viewed 172

I understand that similar questions have been posted here before, but I can't seem to wrap my head around how np.lib_stride_tricks.as_strided works. I would like to know if there is a numpy function that lets me do the following mapping from a -> b efficiently:

import numpy as np

a = np.random.rand(100000,3)

window_length = 200;
spacing = 20;

b = []

for i in range(0,a.shape[0] - window_length,spacing):
    b.append(a[i:i+window_length,:])

b = np.asarray(b);
print(b.shape)

Which outputs (4490,200,3).

3 Answers

We can leverage np.lib.stride_tricks.as_strided based scikit-image's view_as_windows to get sliding windowed views and thus solve our problem here in a vectorized way.

from skimage.util.shape import view_as_windows
w = view_as_windows(a, (window_length,1), step=(spacing,1))[...,0]

# If needed, get into desired permuted axes format
out = w.swapaxes(1,2)

More info on use of as_strided based view_as_windows. If scikit-image is not available, we can directly use the source code that works standalone.

A direct use of as_strided:

strides = a.strides

new_len = (a.shape[0]-window_length+1)//spacing

b = as_strided(a,shape=(new_len,window_length, a.shape[1]),
               strides=(strides[0]*spacing,) + strides)

To generate your array using stride_tricks proceed as follows:

from numpy.lib.stride_tricks import as_strided as stride

d0, d1 = a.shape
s0, s1 = a.strides
b2 = stride(a,
    ((d0 - window_length + 1) // spacing, window_length, d1),
    (s0 * spacing, s0, s1))

I deliberately chose a different array name (b2).

To check that my code yields proper result:

  • generate b array using your code,
  • compare both arrays running np.array_equal(b, b2).

The result should be True.

Related