python - 快速 numpy 滚动

Question

我有一个二维 numpy 数组，我想以增量方式滚动每一行。我在 for 循环中使用 np.roll 来执行此操作。但是因为我调用了数千次，所以我的代码真的很慢。你能帮我看看如何让它更快吗？

我的输入看起来像

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

我的输出看起来像

array([[4,1],
       [2,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]])

既然您已经提到要多次调用这样的滚动例程，请创建一次索引数组 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 的改进!