如何在python中快速循环矩阵计算答案

【问题标题】：How to make fast looping for matrix calculation in python如何在python中快速循环矩阵计算
【发布时间】：2020-07-02 13:01:47
【问题描述】：

我的问题是这个 for 循环需要很长时间才能完成。我想要一种更快的方法来完成它。我的代码是：

dx = 20
dy = 20
dz = 20

x = np.arange(0, 1201, dx)
y = np.arange(0, 1001, dy)
z = np.arange(20, 501, dz)
drho = 3000  # Delta Rho (Density Contrast) kg/m^3

# Input Rho to Model
M = np.zeros((len(z), len(x), len(y)))
M[6:16, 26:36, 15:25] = drho
m = np.array(M.flat)
# p (61, 1525)
# M(25, 61)
#  m(1525,)
# Station Position
stx, sty = np.meshgrid(x, y)
stx = np.array(stx.flat)
sty = np.array(sty.flat)
stz = np.zeros(len(stx))

# Make meshgrid
X, Y, Z = np.meshgrid(x, y, z)
X = np.array(X.flat)
Y = np.array(Y.flat)
Z = np.array(Z.flat)

p = np.zeros((len(stx), len(X)))

# p(3111, 77775)
for i in range(len(X)):
    for j in range(len(stx)):
        p[j, i] = (Z[i] - stz[j]) / ((Z[i] - stz[j]) ** 2 + (X[i] - stx[j]) ** 2 + (Y[i] - sty[j]) ** 2) ** (3/2)

迭代变量有时会超过 2.41 亿次，而且会一直持续下去。

【问题讨论】：

标签： python performance numpy for-loop iteration

【解决方案1】：

您可以通过使用broadcasting 来提高性能：

p = (Z - stz[:,None]) / ((Z - stz[:,None])**2  + (X - stx[:,None])**2 + (Y - sty[:,None])**2) ** (3/2)

请注意，正如 jerome 所指出的那样，这里的性能改进将以内存效率为代价。

检查和计时使用：

x = np.arange(0, 121, dx)
y = np.arange(0, 101, dy)
z = np.arange(20, 51, dz)

def op():
    p = np.zeros((len(stx), len(X)))
    for i in range(len(X)):
        for j in range(len(stx)):
            p[j, i] = (Z[i] - stz[j]) / ((Z[i] - stz[j]) ** 2 + (X[i] - stx[j]) ** 2 + (Y[i] - sty[j]) ** 2) ** (3/2)
    return p

def ap_1():
    return (Z - stz[:,None]) / ((Z - stz[:,None])**2  + (X - stx[:,None])**2 + (Y - sty[:,None])**2) ** (3/2)

%timeit p = op()
# 44.5 ms ± 3.58 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit p_ = ap_1()
# 169 µs ± 2.27 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

np.allclose(p_, p)
# True

我们可以通过让numexpr 处理算术来进一步提升性能并提高内存效率：

import numexpr as ne

def ap_2():
    return ne.evaluate('(Z - stz2D) / ((Z - stz2D)**2  + (X - stx2D)**2 + (Y - sty2D)**2) ** (3/2)',
           {'stz2D':stz[:,None], 'stx2D':stx[:,None], 'sty2D':sty[:,None]})

%timeit ap_2()
# 106 µs ± 6.34 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

因此，通过第二种方法，我们获得了 420x 加速

【讨论】：

请注意，第一个解决方案比第二个解决方案占用更多内存（在我的机器上，第一个解决方案大约 7 Go，第二个解决方案大约 1.8 Go）。
是的，这将以牺牲内存效率为代价，答案@jerome

【解决方案2】：

python 的一个非常酷的地方是 map()、filter() 和 reduce() 函数针对大型数据集进行了高度优化。

如果您将 for 循环替换为这些函数，您应该会看到性能略有提升。（或者一个大的..也许）。

【讨论】：