Question

我需要加快一个for循环，它会执行类似下面代码的操作：

import numpy as np

x = np.random.normal(size=(206,11,11))
y = np.random.normal(size=(206,11,11))
complx = x + 1j*y
complx

a,b,c = complx.shape

for n in xrange(a):
    #do somthing
    z = np.zeros(b) 
    for i in xrange(b):       
        z[i] = (complx[n,:,:].real[i][i]*complx[n,:,:].real[i][i] +\
                complx[n,:,:].imag[i][i]*complx[n,:,:].imag[i][i])(**-0.25)

我模糊地意识到这些事情有时可以用numpy.einsum完成。

但是，我不确定如何使用它？

或者有人有任何其他建议吗？

Answer 1

如果你想加速内部for循环，你可以做类似的事情

import numpy as np

x = np.random.normal(size=(206,11,11))
y = np.random.normal(size=(206,11,11))
complx = x + 1j*y

# takes only the diagonal part of all the 11x11 matrices
complx_diag = np.diagonal(complx,0,1,2)

# do the calc
zn = np.abs(complx_diag)**(-0.5)

for n in xrange(a):
    z = zn[n]
    # do your stuff

如果您的stuff不太复杂，也可以进行矢量化（非常可能）。在for循环之外计算得越多，代码就越快。

Answer 2

如果我没有记错的话，这或多或少都是你想要的。打印陈述只是为了说服自己计算是正确的。

def optimize_01():
    x = np.random.normal(size=(6, 11, 11))
    y = np.random.normal(size=(6, 11, 11))
    complx = x + 1j * y

    a, b, _ = complx.shape

    for n in range(a):
        # do somthing
        A = complx[n, :, :]
        d = np.diagonal(A)
        z = np.power(np.abs(d * d), -0.25)

        print (d[0])
        print (z[0])
        print ((d[0].real * d[0].real + d[0].imag * d[0].imag) ** -0.25)

编辑：如果我将此实现与您的实现进行比较，我会得到以下结果。

import timeit

def optimize_02():
    x = np.random.normal(size=(206, 11, 11))
    y = np.random.normal(size=(206, 11, 11))
    complx = x + 1j * y

    a, b, _ = complx.shape

    for n in range(a):
        # do somthing
        A = complx[n, :, :]
        d = np.diagonal(A)
        z = np.power(np.abs(d * d), -0.25)

def optimize_03():
    x = np.random.normal(size=(206, 11, 11))
    y = np.random.normal(size=(206, 11, 11))
    complx = x + 1j * y

    a, b, _ = complx.shape

    for n in range(a):
        # do somthing
        z = np.zeros(b) 
        for i in range(b):       
            z[i] = (complx[n, :, :].real[i][i] * complx[n, :, :].real[i][i] + \
                    complx[n, :, :].imag[i][i] * complx[n, :, :].imag[i][i]) ** (-0.25)

if __name__ == '__main__':
    print (timeit.timeit(optimize_02, number=10))
    print (timeit.timeit(optimize_03, number=10))

结果：

0.03474012700007734
0.09025639800074714

使用6个1100个元素的数组，而不是206个11个元素的数组，结果是：

5.762741210999593
5.771216576999905

看起来我的解决方案毕竟不是那么快。

加速三轴循环＃39;在NumPy

2 个答案: