Hey,
Inspired by an existing PR into numpy, I created two functions based on
stride_tricks which I thought might be useful inside numpy. So if I get
any feedback/pointers, I would add some tests and create a PR.
The first function rolling_window is to create for every point of the
original array, a view of its neighbourhood in arbitrary dimensions with
some options to make it easy to use.
For example for a 3 dim array:
rolling_window(a, (3,3,3), asteps=(3,3,3)).max((3,4,5))
Gives the maximum for all non-overlapping 3x3x3 subarrays and:
rolling_window(a, 3, axes=0).max(-1) would create the rolling maximum
over all 3 successive values along the 0th axis.
Plus a function permute_axes which allows to give a tuple for switching
axes and can also merge them. A (10, 4, 3) shaped array would give:
permute_axes(a, (2, 1, 0)).shape -> (3, 4, 10)
permute_axes(a, (1, (0, 2))).shape -> (4, 30)
A copy is created when necessary, this basically allows to do multiple
swap_axes and a reshape in a single call.
I hope this might be useful...
Regards,
Sebastian
import numpy as np
def rolling_window(array, window=(0,), asteps=None, wsteps=None, axes=None, intersperse=False):
"""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.
intersperse : bool
If True, the new dimensions are right after the corresponding original
dimension, instead of at the end of the 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."
