快速numpy卷

Question

我有一個2D numpy陣列，我想以增量方式滾動每一行。我在for循環中使用np.roll來執行此操作。但是因爲我打了幾千次這樣的代碼，所以我的代碼非常慢。你能幫我解決一下如何讓它更快。

我輸入看起來像

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

和我的輸出看起來像

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

這裏的0行[4,1]由0偏移，並且第一排[0,2]被1.類似地，第二移行將被移動2，依此類推。

編輯

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

Answer 1

這裏有一個量化的解決方案 -

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

樣品輸入，輸出 -

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

因爲，你所提到的，你在呼喚這樣的壓延常規多次，創建索引陣列y idx一次，稍後重新使用。

進一步改善

重複用法，你是最好創建完整的線性指標，然後使用np.take提取捲起的元素，就像這樣 -

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

讓我們來看看有什麼改進像 -

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 

大約3x有改進！

Answer 2

一個棘手，但快速的解決方案：

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,(m-n,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]]