chrisaycock chrisaycock - 26 days ago 8
Python Question

Shift elements in a numpy array

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

Answer

Not numpy but scipy provides exactly the shift functionality you want,

import numpy as np
from scipy.ndimage.interpolation import shift

xs = np.array([ 0.,  1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9.])

shift(xs, 3, cval=np.NaN)

where default is to bring in a constant value from outside the array with value cval, set here to nan. This gives the desired output,

array([ nan, nan, nan, 0., 1., 2., 3., 4., 5., 6.])

and the negative shift works similarly,

shift(xs, -3, cval=np.NaN)

Provides output

array([  3.,   4.,   5.,   6.,   7.,   8.,   9.,  nan,  nan,  nan])