Shift Elements in a Numpy Array

Question

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.

Accepted 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])

Accepted Answer

For those who want to just copy and paste the fastest implementation of shift, there is a benchmark and conclusion(see the end). In addition, I introduce fill_value parameter and fix some bugs.

Benchmark

import numpy as np
import timeit

# enhanced from IronManMark20 version
def shift1(arr, num, fill_value=np.nan):
    arr = np.roll(arr,num)
    if num < 0:
        arr[num:] = fill_value
    elif num > 0:
        arr[:num] = fill_value
    return arr

# use np.roll and np.put by IronManMark20
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

# use np.pad and slice by me.
def shift3(arr, num, fill_value=np.nan):
    l = len(arr)
    if num < 0:
        arr = np.pad(arr, (0, abs(num)), mode='constant', constant_values=(fill_value,))[:-num]
    elif num > 0:
        arr = np.pad(arr, (num, 0), mode='constant', constant_values=(fill_value,))[:-num]

    return arr

# use np.concatenate and np.full by chrisaycock
def shift4(arr, num, fill_value=np.nan):
    if num >= 0:
        return np.concatenate((np.full(num, fill_value), arr[:-num]))
    else:
        return np.concatenate((arr[-num:], np.full(-num, fill_value)))

# preallocate empty array and assign slice by chrisaycock
def shift5(arr, num, fill_value=np.nan):
    result = np.empty_like(arr)
    if num > 0:
        result[:num] = fill_value
        result[num:] = arr[:-num]
    elif num < 0:
        result[num:] = fill_value
        result[:num] = arr[-num:]
    else:
        result[:] = arr
    return result

arr = np.arange(2000).astype(float)

def benchmark_shift1():
    shift1(arr, 3)

def benchmark_shift2():
    shift2(arr, 3)

def benchmark_shift3():
    shift3(arr, 3)

def benchmark_shift4():
    shift4(arr, 3)

def benchmark_shift5():
    shift5(arr, 3)

benchmark_set = ['benchmark_shift1', 'benchmark_shift2', 'benchmark_shift3', 'benchmark_shift4', 'benchmark_shift5']

for x in benchmark_set:
    number = 10000
    t = timeit.timeit('%s()' % x, 'from __main__ import %s' % x, number=number)
    print '%s time: %f' % (x, t)

benchmark result:

benchmark_shift1 time: 0.265238
benchmark_shift2 time: 0.285175
benchmark_shift3 time: 0.473890
benchmark_shift4 time: 0.099049
benchmark_shift5 time: 0.052836

Conclusion

shift5 is winner! It's OP's third solution.

Answer 3

Benchmarks & introducing Numba

Must Read

1. Summary

The accepted answer (scipy.ndimage.interpolation.shift) is the slowest solution listed in this page.
Numba (@numba.njit) gives some performance boost when array size smaller than ~25.000
"Any method" equally good when array size large (>250.000).
The fastest option really depends on
(1) Length of your arrays
(2) Amount of shift you need to do.
Below is the picture of the timings of all different methods listed on this page (2020-07-11), using constant shift = 10. As one can see, with small array sizes some methods are use more than +2000% time than the best method.