如何在python中進行矩陣計算的快速循環

Question

我有一個問題，這個 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 億，並且永遠如此。

Answer 1

您可以通過使用 廣播來提高性能：

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加速

Answer 2

關於 python 的一個非常酷的事情是 map()、filter() 和 reduce() 函數針對大型數據集進行了大量優化。

如果你用這些函數替換了你的for循環，你應該會看到一個小的性能改進。 （或者一個大的......也許）。