Per line numpy roll

Question

Is there an elegant way to perform a per-line roll of a numpy array? Example:

>>> arr
array([[1, 2, 3, 4],
       [5, 6, 7, 8]])
>>> desired_func(arr, shift=np.array([1, 2]))
array([[4, 1, 2, 3],
       [7, 8, 5, 6]])

The function np.roll doesn't seem to allow that.

Answer 1

You can use fancy indexing:

# define the roll per row
roll = np.array([1, 2])

# compute the rolled indices 
idx = (np.arange(arr.shape[1]) - roll[:,None]) % arr.shape[1]

# index as 2D
out = arr[np.arange(arr.shape[0])[:,None], idx]

output:

array([[4, 1, 2, 3],
       [7, 8, 5, 6]])