Skip to content

Instantly share code, notes, and snippets.

@seberg
Created October 10, 2012 14:38
Show Gist options
  • Save seberg/3866040 to your computer and use it in GitHub Desktop.
Save seberg/3866040 to your computer and use it in GitHub Desktop.
Multidimensional rolling_window for numpy
def rolling_window(array, window=(0,), asteps=None, wsteps=None, axes=None, toend=True):
"""Create a view of `array` which for every point gives the n-dimensional
neighbourhood of size window. New dimensions are added at the end of
`array` or after the corresponding original dimension.
Parameters
----------
array : array_like
Array to which the rolling window is applied.
window : int or tuple
Either a single integer to create a window of only the last axis or a
tuple to create it for the last len(window) axes. 0 can be used as a
to ignore a dimension in the window.
asteps : tuple
Aligned at the last axis, new steps for the original array, ie. for
creation of non-overlapping windows. (Equivalent to slicing result)
wsteps : int or tuple (same size as window)
steps for the added window dimensions. These can be 0 to repeat values
along the axis.
axes: int or tuple
If given, must have the same size as window. In this case window is
interpreted as the size in the dimension given by axes. IE. a window
of (2, 1) is equivalent to window=2 and axis=-2.
toend : bool
If False, the new dimensions are right after the corresponding original
dimension, instead of at the end of the array. Adding the new axes at the
end makes it easier to get the neighborhood, however toend=False will give
a more intuitive result if you view the whole array.
Returns
-------
A view on `array` which is smaller to fit the windows and has windows added
dimensions (0s not counting), ie. every point of `array` is an array of size
window.
Examples
--------
>>> a = np.arange(9).reshape(3,3)
>>> rolling_window(a, (2,2))
array([[[[0, 1],
[3, 4]],
[[1, 2],
[4, 5]]],
[[[3, 4],
[6, 7]],
[[4, 5],
[7, 8]]]])
Or to create non-overlapping windows, but only along the first dimension:
>>> rolling_window(a, (2,0), asteps=(2,1))
array([[[0, 3],
[1, 4],
[2, 5]]])
Note that the 0 is discared, so that the output dimension is 3:
>>> rolling_window(a, (2,0), asteps=(2,1)).shape
(1, 3, 2)
This is useful for example to calculate the maximum in all (overlapping)
2x2 submatrixes:
>>> rolling_window(a, (2,2)).max((2,3))
array([[4, 5],
[7, 8]])
Or delay embedding (3D embedding with delay 2):
>>> x = np.arange(10)
>>> rolling_window(x, 3, wsteps=2)
array([[0, 2, 4],
[1, 3, 5],
[2, 4, 6],
[3, 5, 7],
[4, 6, 8],
[5, 7, 9]])
"""
array = np.asarray(array)
orig_shape = np.asarray(array.shape)
window = np.atleast_1d(window).astype(int) # maybe crude to cast to int...
if axes is not None:
axes = np.atleast_1d(axes)
w = np.zeros(array.ndim, dtype=int)
for axis, size in zip(axes, window):
w[axis] = size
window = w
# Check if window is legal:
if window.ndim > 1:
raise ValueError("`window` must be one-dimensional.")
if np.any(window < 0):
raise ValueError("All elements of `window` must be larger then 1.")
if len(array.shape) < len(window):
raise ValueError("`window` length must be less or equal `array` dimension.")
_asteps = np.ones_like(orig_shape)
if asteps is not None:
asteps = np.atleast_1d(asteps)
if asteps.ndim != 1:
raise ValueError("`asteps` must be either a scalar or one dimensional.")
if len(asteps) > array.ndim:
raise ValueError("`asteps` cannot be longer then the `array` dimension.")
# does not enforce alignment, so that steps can be same as window too.
_asteps[-len(asteps):] = asteps
if np.any(asteps < 1):
raise ValueError("All elements of `asteps` must be larger then 1.")
asteps = _asteps
_wsteps = np.ones_like(window)
if wsteps is not None:
wsteps = np.atleast_1d(wsteps)
if wsteps.shape != window.shape:
raise ValueError("`wsteps` must have the same shape as `window`.")
if np.any(wsteps < 0):
raise ValueError("All elements of `wsteps` must be larger then 0.")
_wsteps[:] = wsteps
_wsteps[window == 0] = 1 # make sure that steps are 1 for non-existing dims.
wsteps = _wsteps
# Check that the window would not be larger then the original:
if np.any(orig_shape[-len(window):] < window * wsteps):
raise ValueError("`window` * `wsteps` larger then `array` in at least one dimension.")
new_shape = orig_shape # just renaming...
# For calculating the new shape 0s must act like 1s:
_window = window.copy()
_window[_window==0] = 1
new_shape[-len(window):] += wsteps - _window * wsteps
new_shape = (new_shape + asteps - 1) // asteps
# make sure the new_shape is at least 1 in any "old" dimension (ie. steps
# is (too) large, but we do not care.
new_shape[new_shape < 1] = 1
shape = new_shape
strides = np.asarray(array.strides)
strides *= asteps
new_strides = array.strides[-len(window):] * wsteps
# The full new shape and strides:
if toend:
new_shape = np.concatenate((shape, window))
new_strides = np.concatenate((strides, new_strides))
else:
_ = np.zeros_like(shape)
_[-len(window):] = window
_window = _.copy()
_[-len(window):] = new_strides
_new_strides = _
new_shape = np.zeros(len(shape)*2, dtype=int)
new_strides = np.zeros(len(shape)*2, dtype=int)
new_shape[::2] = shape
new_strides[::2] = strides
new_shape[1::2] = _window
new_strides[1::2] = _new_strides
new_strides = new_strides[new_shape != 0]
new_shape = new_shape[new_shape != 0]
return np.lib.stride_tricks.as_strided(array, shape=new_shape, strides=new_strides)
@mada0007
Copy link

mada0007 commented Feb 2, 2019

Hi Serberg and all following this thread, So I have a gridded data of shape (324,72,144) in order time,lat, and lon. I want to apply this rolling function this way, from time 1 to 60, count the number of months above certain threshold, repeat this for time 2 to 61, 3 t0 62 and so on. So my window is 60 and step size is one but this should be done on the time axis. Can you please explain how I could do this using this function or any other function you might deem fit for use. Most appreciated,

Mustapha.

@LuizFonseca06
Copy link

Thank you!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment