[Numpy-discussion] Rolling window (moving average, moving std, and more)
Sebastian Haase
seb.haase at gmail.com
Mon Jan 3 05:13:25 EST 2011
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
More information about the NumPy-Discussion
mailing list