How To Randomly Shift Rows Of A Numpy Array

I am looking for a more pythonic way of randomly shifting rows of a numpy array. The idea is that I have an array of data, and I want to left-shift each row of the array by a rando

Solution 1:

Generate all the column indices for all rows in one go and then simply use integer-indexing for a vectorized solution, like so -

# Store shape of input array
m,n = data.shape

# Get random column start indices for each row in one go
col_start = np.random.randint(0, max_shift, data.shape[0])

# Get the rolled indices for every row again in a vectorized manner.# We are extending col_start to 2D and then adding a range array to get # all column indices for every row by leveraging NumPy's braodcasting.# Because of the additions, we might go off-limits. So, to simulate the # rolled over version, mod it.
idx = np.mod(col_start[:,None] + np.arange(n), n)

# Finall with integer indexing get the values off data array
shifted_out = data[np.arange(m)[:,None], idx]

Step-by-step run -

1] Inputs :

In [548]: data
Out[548]: 
array([[44, 23, 38, 32, 30],
       [69, 15, 32, 41, 63],
       [69, 41, 75, 50, 87],
       [23, 28, 38, 79, 91]])

In [549]: max_shift = 5

2] Proposed solution :

2A] Get column starts :

In [550]: m,n = data.shape

In [551]: col_start = np.random.randint(0, max_shift, data.shape[0])

In [552]: col_start
Out[552]: array([1, 2, 3, 3])

2B] Get all indices :

In [553]: idx = np.mod(col_start[:,None] + np.arange(n), n)

In [554]: col_start[:,None]
Out[554]: 
array([[1],
       [2],
       [3],
       [3]])

In [555]: col_start[:,None] + np.arange(n)
Out[555]: 
array([[1, 2, 3, 4, 5],
       [2, 3, 4, 5, 6],
       [3, 4, 5, 6, 7],
       [3, 4, 5, 6, 7]])

In [556]: np.mod(col_start[:,None] + np.arange(n), n)
Out[556]: 
array([[1, 2, 3, 4, 0],
       [2, 3, 4, 0, 1],
       [3, 4, 0, 1, 2],
       [3, 4, 0, 1, 2]])

2C] Finally index into data :

In [557]: data[np.arange(m)[:,None], idx]
Out[557]: 
array([[23, 38, 32, 30, 44],
       [32, 41, 63, 69, 15],
       [50, 87, 69, 41, 75],
       [79, 91, 23, 28, 38]])

Verification -

1] Original approach :

In [536]: data = np.random.randint(11,99,(4,5))
     ...: max_shift = 5
     ...: col_start = -np.random.randint(0, max_shift, data.shape[0])
     ...: for i,row in enumerate(data):
     ...:     print np.array([np.roll(row, col_start[i])])
     ...:     
[[83 93 17 53 61]][[55 88 84 94 89]][[59 63 29 72 85]][[57 95 13 21 14]]

2] Proposed approach re-using col_start, so that we could do a value verification :

In [537]: m,n = data.shape

In [538]: idx = np.mod(-col_start[:,None] + np.arange(n), n)

In [539]: data[np.arange(m)[:,None], idx]
Out[539]: 
array([[83, 93, 17, 53, 61],
       [55, 88, 84, 94, 89],
       [59, 63, 29, 72, 85],
       [57, 95, 13, 21, 14]])