加速 NumPy 中三个轴的循环

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

如果我没记错的话，这或多或少就是你想要的。打印语句只是为了让自己相信计算是正确的。

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

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

5.762741210999593
5.771216576999905

看来我的解决方案毕竟没有那么快。

Answer 2

如果你想加速内部 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 循环外计算得越多，代码运行得越快。

加速 NumPy 中三个轴的循环

Speed up looping over three axis' in NumPy

python

numpy