Quick roll

Question

I have a 2d numpy array and I want to rotate each line in stages. For this, I use np.roll in a for loop. But since I call a thousand times, my code is very slow. Could you help me find out how to do this faster.

My input looks like

 array([[4,1], [0,2]]) 

and my conclusion looks like

 array([[4,1], [2,0]]) 

Here, the zero line [4,1] was shifted by 0, and the first line [0,2] was shifted by 1. Similarly, the second line will be shifted by 2, etc.

EDIT

 temp = np.zeros([dd,dd]) for i in range(min(t + 1, dd)): temp[i,:] = np.roll(y[i,:], i, axis=0) 

Answer 1

Here is one vector solution -

 m,n = a.shape idx = np.mod((n-1)*np.arange(m)[:,None] + np.arange(n), n) out = a[np.arange(m)[:,None], idx] 

Example input, output -

 In [256]: a Out[256]: array([[73, 55, 79, 52, 15], [45, 11, 19, 93, 12], [78, 50, 30, 88, 53], [98, 13, 58, 34, 35]]) In [257]: out Out[257]: array([[73, 55, 79, 52, 15], [12, 45, 11, 19, 93], [88, 53, 78, 50, 30], [58, 34, 35, 98, 13]]) 

Since you mentioned that you call such a looping procedure several times, create an idx indexing array once and reuse it later.

Further improvement

For repeated uses, you better create full linear indexes and then use np.take to retrieve collapsed elements, for example:

 full_idx = idx + n*np.arange(m)[:,None] out = np.take(a,full_idx) 

Let's see what an improvement is -

 In [330]: a = np.random.randint(11,99,(600,600)) In [331]: m,n = a.shape ...: idx = np.mod((n-1)*np.arange(m)[:,None] + np.arange(n), n) ...: In [332]: full_idx = idx + n*np.arange(m)[:,None] In [333]: %timeit a[np.arange(m)[:,None], idx] # Approach #1 1000 loops, best of 3: 1.42 ms per loop In [334]: %timeit np.take(a,full_idx) # Improvement 1000 loops, best of 3: 486 µs per loop 

Around 3x improvement!

Answer 2

difficult but quick solution:

 p=5 a=randint(0,p,(p,p)) aa=hstack((a,a)) m,n=aa.strides b=np.lib.stride_tricks.as_strided(aa,a.shape,(m+n,n)) c=np.lib.stride_tricks.as_strided(aa.ravel()[p:],a.shape,(mn,n)) ## [[2 1 4 2 4] [0 4 2 0 3] [1 3 3 4 4] [1 0 3 2 4] [3 3 2 1 3]] [[2 1 4 2 4] [4 2 0 3 0] [3 4 4 1 3] [2 4 1 0 3] [3 3 3 2 1]] [[2 1 4 2 4] [3 0 4 2 0] [4 4 1 3 3] [3 2 4 1 0] [3 2 1 3 3]]