如何随机移动numpy数组的行

Question

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 random amount. My solution, which works, but I feel is a bit un-pythonic:

def shift_rows(data, max_shift):
    """Left-shifts each row in `data` by a random amount up to `max_shift`."""
    return np.array([np.roll(row, -np.random.randint(0, max_shift)) for row in data])

And to test:

data = np.array([np.arange(0, 5) for _ in range(10)])  # toy data to illustrate
shifted = shift_rows(data, max_shift=5)
shifted
# array([1, 2, 3, 4, 0],
#       [1, 2, 3, 4, 0],
#       [0, 1, 2, 3, 4],
#       ...
#       [4, 0, 1, 2, 3]])

This is really more of a thought experiment. Can anybody come up with a more efficient or more pythonic way of doing this? I suppose list comprehensions are pythonic, but if I need to do this over a huge array is this efficient?

Edit: I marked the excellent reply by Divakar as the answer, but I would still love to hear it if anybody has any other ideas.

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