有效地重塑numpy数组

时间:2016-06-14 19:01:28

标签: python arrays performance numpy vectorization

我正在使用NumPy数组。

我有一个2N长度向量D,想要将其中的一部分重新整形为N x N数组C

现在这段代码符合我的要求,但却是更大N的瓶颈:

```

import numpy as np
M = 1000
t = np.arange(M)
D = np.sin(t)    # initial vector is a sin() function
N = M / 2
C = np.zeros((N,N))
for a in xrange(N):
    for b in xrange(N):
        C[a,b] = D[N + a - b]

```

一旦C生成,我就继续对它进行矩阵运算等等。

这个嵌套循环很慢,但由于这个操作基本上是索引的变化,我认为我可以使用NumPy的内置重塑(numpy.reshape)来加速这部分。

不幸的是,我似乎无法找到改变这些指数的好方法。

加速这部分的任何帮助?

1 个答案:

答案 0 :(得分:3)

您可以使用NumPy broadcasting删除那些嵌套循环 -

C = D[N + np.arange(N)[:,None] - np.arange(N)]

还可以使用np.take替换索引,如此 -

C = np.take(D,N + np.arange(N)[:,None] - np.arange(N))

仔细观察会发现该模式接近toeplitzhankel矩阵。因此,使用这些,我们将有两种方法来解决它,尽管具有与广播相当的加速比。实现看起来像这些 -

from scipy.linalg import toeplitz
from scipy.linalg import hankel

C = toeplitz(D[N:],np.hstack((D[0],D[N-1:0:-1])))
C = hankel(D[1:N+1],D[N:])[:,::-1]

运行时测试

In [230]: M = 1000
     ...: t = np.arange(M)
     ...: D = np.sin(t)    # initial vector is a sin() function
     ...: N = M / 2
     ...: 

In [231]: def org_app(D,N):
     ...:     C = np.zeros((N,N))
     ...:     for a in xrange(N):
     ...:         for b in xrange(N):
     ...:             C[a,b] = D[N + a - b]
     ...:     return C
     ...: 

In [232]: %timeit org_app(D,N)
     ...: %timeit D[N + np.arange(N)[:,None] - np.arange(N)]
     ...: %timeit np.take(D,N + np.arange(N)[:,None] - np.arange(N))
     ...: %timeit toeplitz(D[N:],np.hstack((D[0],D[N-1:0:-1])))
     ...: %timeit hankel(D[1:N+1],D[N:])[:,::-1]
     ...: 
10 loops, best of 3: 83 ms per loop
100 loops, best of 3: 2.82 ms per loop
100 loops, best of 3: 2.84 ms per loop
100 loops, best of 3: 2.95 ms per loop
100 loops, best of 3: 2.93 ms per loop