Rolling matrix rows independently

Question

I have a matrix (2d numpy ndarray, to be precise):

A = np.array([[4, 0, 0], [1, 2, 3], [0, 0, 5]]) 

And I want to roll each row of A independently, according to the roll values ​​in another array:

 r = np.array([2, 0, -1]) 

That is, I want to do this:

 print np.array([np.roll(row, x) for row,x in zip(A, r)]) [[0 0 4] [1 2 3] [0 5 0]] 

Is there any way to do this efficiently? Perhaps using fancy indexing tricks?

Answer 1

Of course, you can do this with advanced indexing, maybe this is the fastest way, probably depends on the size of your array (if your lines are large, this may not be the case):

 rows, column_indices = np.ogrid[:A.shape[0], :A.shape[1]] # Use always a negative shift, so that column_indices are valid. # (could also use module operation) r[r < 0] += A.shape[1] column_indices = column_indices - r[:, np.newaxis] result = A[rows, column_indices] 

Answer 2

numpy.lib.stride_tricks.as_strided stricks (abbreviation for pun intended) again!

Speaking of bizarre indexing tricks, there is the notorious np.lib.stride_tricks.as_strided . The idea / trick would be to get the cut off part, starting from the first column, to the second last and concatenation at the end. This ensures that we can move forward when necessary to use np.lib.stride_tricks.as_strided and thus avoid the need for an actual rollback. This is the whole idea!

Now, in terms of actual implementation, we will use scikit-image view_as_windows to elegantly use np.lib.stride_tricks.as_strided under the hoods. So the final implementation will be -

 from skimage.util.shape import view_as_windows as viewW def strided_indexing_roll(a, r): # Concatenate with sliced to cover all rolls a_ext = np.concatenate((a,a[:,:-1]),axis=1) # Get sliding windows; use advanced-indexing to select appropriate ones n = a.shape[1] return viewW(a_ext,(1,n))[np.arange(len(r)), (nr)%n,0] 

Here’s a run sample -

 In [327]: A = np.array([[4, 0, 0], ...: [1, 2, 3], ...: [0, 0, 5]]) In [328]: r = np.array([2, 0, -1]) In [329]: strided_indexing_roll(A, r) Out[329]: array([[0, 0, 4], [1, 2, 3], [0, 5, 0]]) 

Benchmarking

 # @seberg solution def advindexing_roll(A, r): rows, column_indices = np.ogrid[:A.shape[0], :A.shape[1]] r[r < 0] += A.shape[1] column_indices = column_indices - r[:,np.newaxis] return A[rows, column_indices] 

Let me do some benchmarking on an array with lots of rows and columns -

 In [324]: np.random.seed(0) ...: a = np.random.rand(10000,1000) ...: r = np.random.randint(-1000,1000,(10000)) # @seberg solution In [325]: %timeit advindexing_roll(a, r) 10 loops, best of 3: 71.3 ms per loop # Solution from this post In [326]: %timeit strided_indexing_roll(a, r) 10 loops, best of 3: 44 ms per loop 

Answer 3

If you need a more general solution (with any shape and any axis), I changed @seberg's solution:

 def indep_roll(arr, shifts, axis=1): """Apply an independent roll for each dimensions of a single axis. Parameters ---------- arr : np.ndarray Array of any shape. shifts : np.ndarray How many shifting to use for each dimension. Shape: '(arr.shape[axis],)'. axis : int Axis along which elements are shifted. """ arr = np.swapaxes(arr,axis,-1) all_idcs = np.ogrid[[slice(0,n) for n in arr.shape]] # Convert to a positive shift shifts[shifts < 0] += arr.shape[-1] all_idcs[-1] = all_idcs[-1] - shifts[:, np.newaxis] result = arr[tuple(all_idcs)] arr = np.swapaxes(result,-1,axis) return arr