Shift elements in a numpy array

Question

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 o

Answer 1

There is no single function that does what you want. Your definition of shift is slightly different than what most people are doing. The ways to shift an array are more commonly looped:

>>>xs=np.array([1,2,3,4,5])
>>>shift(xs,3)
array([3,4,5,1,2])

However, you can do what you want with two functions.
Consider a=np.array([ 0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]):

def shift2(arr,num):
    arr=np.roll(arr,num)
    if num<0:
         np.put(arr,range(len(arr)+num,len(arr)),np.nan)
    elif num > 0:
         np.put(arr,range(num),np.nan)
    return arr
>>>shift2(a,3)
[ nan  nan  nan   0.   1.   2.   3.   4.   5.   6.]
>>>shift2(a,-3)
[  3.   4.   5.   6.   7.   8.   9.  nan  nan  nan]

After running cProfile on your given function and the above code you provided, I found that the code you provided makes 42 function calls while shift2 made 14 calls when arr is positive and 16 when it is negative. I will be experimenting with timing to see how each performs with real data.