Если вы знаете, что хотите двумерное окно, укажите window_size[0]
для array.shape[0]
и window_size[1]
для array.shape[1]
при определении shape
:
def rolling_window(array, window_size):
itemsize = array.itemsize
# broadcast window_size so it will still work with a scalar:
window_size = np.broadcast_to(window_size,[2,])
shape = (array.shape[0] - window_size[0] + 1,
array.shape[1] - window_size[1] + 1,
window_size[0], window_size[1])
# the following would also work:
# shape = (*np.array(array.shape) - window_size + 1, *window_size)
strides = (array.shape[1] * itemsize, itemsize,
array.shape[1] * itemsize, itemsize)
return np.lib.stride_tricks.as_strided(array, shape=shape, strides=strides)
roller = np.arange(1,16).reshape(3,5)
>>> print(roller)
[[ 1 2 3 4 5]
[ 6 7 8 9 10]
[11 12 13 14 15]]
run = rolling_window(roller,[2,3])
>>> run
array([[[[ 1, 2, 3],
[ 6, 7, 8]],
[[ 2, 3, 4],
[ 7, 8, 9]],
[[ 3, 4, 5],
[ 8, 9, 10]]],
[[[ 6, 7, 8],
[11, 12, 13]],
[[ 7, 8, 9],
[12, 13, 14]],
[[ 8, 9, 10],
[13, 14, 15]]]])