numpy.roll horizontally on a 2D ndarray with different values

Question

Doing np.roll(a, 1, axis = 1) on:

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

results in the correct:

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

The documentation says:

If a tuple, then axis must be a tuple of the same size, and each of the given axes is shifted by the corresponding number.

Now I like to roll rows of a by different values, like [1,2,1,3] meaning, first row will be rolled by 1, second by 2, third by 1 and forth by 3. But np.roll(a, [1,2,1,3], axis=(1,1,1,1)) doesn't seem to do it. What would be the correct interpretation of the sentence in the docs?

Answer 1

By specifying a tuple in np.roll you can roll an array along various axes. For example, np.roll(a, (3,2), axis=(0,1)) will shift each element of a by 3 places along axis 0, and it will also shift each element by 2 places along axis 1. np.roll does not have an option to roll each row by a different amount. You can do it though for example as follows:

import numpy as np

a = np.array([  
             [6, 3, 9, 2, 3],
             [1, 7, 8, 1, 2],
             [5, 4, 2, 2, 4],
             [3, 9, 7, 6, 5],
            ])
shifts = np.c_[[1,2,1,3]]
a[np.c_[:a.shape[0]], (np.r_[:a.shape[1]] - shifts) % a.shape[1]]

It gives:

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