[Python] Shift elements in a numpy array

This question contains its own answer at the bottom. Use preallocated arrays.

Following-up from this question years ago, is there a canonical "shift" function in numpy? I don't see anything from the documentation.

Here's a simple version of what I'm looking for:

def shift(xs, n):
if n >= 0:
return np.r_[np.full(n, np.nan), xs[:-n]]
else:
return np.r_[xs[-n:], np.full(-n, np.nan)]


Using this is like:

In [76]: xs
Out[76]: array([ 0., 1., 2., 3., 4., 5., 6., 7., 8., 9.])

In [77]: shift(xs, 3)
Out[77]: array([ nan, nan, nan, 0., 1., 2., 3., 4., 5., 6.])

In [78]: shift(xs, -3)
Out[78]: array([ 3., 4., 5., 6., 7., 8., 9., nan, nan, nan])


This question came from my attempt to write a fast rolling_product yesterday. I needed a way to "shift" a cumulative product and all I could think of was to replicate the logic in np.roll().


So np.concatenate() is much faster than np.r_[]. This version of the function performs a lot better:

def shift(xs, n):
if n >= 0:
return np.concatenate((np.full(n, np.nan), xs[:-n]))
else:
return np.concatenate((xs[-n:], np.full(-n, np.nan)))


An even faster version simply pre-allocates the array:

def shift(xs, n):
e = np.empty_like(xs)
if n >= 0:
e[:n] = np.nan
e[n:] = xs[:-n]
else:
e[n:] = np.nan
e[:n] = xs[-n:]
return e


The above proposal is the answer. Use preallocated arrays.

Continue reading...