15 lines
539 B
Python
15 lines
539 B
Python
import numpy as np
|
|
|
|
def rolling(a, window):
|
|
# http://stackoverflow.com/a/4924433
|
|
shape = (a.size - window + 1, window)
|
|
strides = (a.itemsize, a.itemsize)
|
|
return np.lib.stride_tricks.as_strided(a, shape=shape, strides=strides)
|
|
|
|
def rolling_batch(a, window):
|
|
# same as rolling, but acts on each batch (axis 0).
|
|
shape = (a.shape[0], a.shape[-1] - window + 1, window)
|
|
strides = (np.prod(a.shape[1:]) * a.itemsize, a.itemsize, a.itemsize)
|
|
return np.lib.stride_tricks.as_strided(a, shape=shape, strides=strides)
|
|
|