[Numpy-discussion] Rolling window (moving average, moving std, and more)

[Sebastian Haase]

Hi Erik,
This is really neat !  Do I understand correctly, that you mean by
"stride tricks", that your rolling_window is _not_ allocating any new
memory ?
IOW, If I have a large array using 500MB of memory, say of float32 of
shape 125,1000,1000 and I want the last axis rolling of window size
11, what would the peak memory usage of that operation be ?
How about renaming the option `window` to `window_size`  (first I was
thinking of things like hamming and hanning windows...)... ?

Thanks,
Sebastian Haase


On Sat, Jan 1, 2011 at 5:29 AM, Erik Rigtorp <erik at rigtorp.com> wrote:
> Hi,
>
> Implementing moving average, moving std and other functions working
> over rolling windows using python for loops are slow. This is a
> effective stride trick I learned from Keith Goodman's
> <kwgoodman at gmail.com> Bottleneck code but generalized into arrays of
> any dimension. This trick allows the loop to be performed in C code
> and in the future hopefully using multiple cores.
>
> import numpy as np
>
> def rolling_window(a, window):
>    """
>    Make an ndarray with a rolling window of the last dimension
>
>    Parameters
>    ----------
>    a : array_like
>        Array to add rolling window to
>    window : int
>        Size of rolling window
>
>    Returns
>    -------
>    Array that is a view of the original array with a added dimension
>    of size w.
>
>    Examples
>    --------
>    >>> x=np.arange(10).reshape((2,5))
>    >>> rolling_window(x, 3)
>    array([[[0, 1, 2], [1, 2, 3], [2, 3, 4]],
>           [[5, 6, 7], [6, 7, 8], [7, 8, 9]]])
>
>    Calculate rolling mean of last dimension:
>    >>> np.mean(rolling_window(x, 3), -1)
>    array([[ 1.,  2.,  3.],
>           [ 6.,  7.,  8.]])
>
>    """
>    if window < 1:
>        raise ValueError, "`window` must be at least 1."
>    if window > a.shape[-1]:
>        raise ValueError, "`window` is too long."
>    shape = a.shape[:-1] + (a.shape[-1] - window + 1, window)
>    strides = a.strides + (a.strides[-1],)
>    return np.lib.stride_tricks.as_strided(a, shape=shape, strides=strides)
>
>
> Using np.swapaxes(-1, axis) rolling aggregations over any axis can be computed.
>
> I submitted a pull request to add this to the stride_tricks module.
>
> Erik