rougier/numpy-100

Solution for 76 (Can be wrong when use a view of array as input)

xzymustbexzy opened this issue · 2 comments

76. Consider a one-dimensional array Z, build a two-dimensional array whose first row is (Z[0],Z[1],Z[2]) and each subsequent row is shifted by 1 (last row should be (Z[-3],Z[-2],Z[-1])

# Author: Joe Kington / Erik Rigtorp
from numpy.lib import stride_tricks

def rolling(a, window):
    shape = (a.size - window + 1, window)
    strides = (a.itemsize, a.itemsize)
    return stride_tricks.as_strided(a, shape=shape, strides=strides)

Z = rolling(np.arange(10), 3)
print(Z)

When use view of array with strides as input.

a = np.arange(10)
a = a[::2]
Z = rolling(a, 3)
print(Z)

The output

[[0 1 2]
 [1 2 3]
 [2 3 4]]

My solution is using strides instead of itemsize. Therefore, the answer looks like this:

from numpy.lib import stride_tricks

def rolling(a, window):
    shape = (a.size - window + 1, window)
    strides = (a.strides[0], a.strides[0])
    return stride_tricks.as_strided(a, shape=shape, strides=strides)

In this way, the result is correct.

[[0 2 4]
 [2 4 6]
 [4 6 8]]

Good catch! This is actually an error I think, thanks. Can you make a PR?

Good catch! This is actually an error I think, thanks. Can you make a PR?

Okay!