一位用户在将MATLAB中的高斯-赛德尔(高斯-雅可比)求解器代码移植到NumPy后,发现NumPy版本的求解速度比MATLAB慢了近1.5倍。该代码用于解决二维轴对称域(柱坐标系)中的离散拉普拉斯方程,并以指定公差停止迭代。用户尝试了将 NumPy 数组更改为列优先顺序,但并没有提高性能。
2、解决方案
- 减少中间数组的创建。
通过减少创建的中间数组的数量,包括重用预先分配的工作数组,可以显著提高求解器的速度。以下是优化后的代码:
import numpy as np
import time
T = np.transpose
# geometry
length = 0.008
width = 0.002
# mesh
nz = 256
nr = 64
# step sizes
dz = length/nz
dr = width/nr
# node position matrices
r = np.tile(np.linspace(0,width,nr+1), (nz+1, 1)).T
ri = r/dr
# equation coefficients
cr = dz**2 / (2*(dr**2 + dz**2))
cz = dr**2 / (2*(dr**2 + dz**2))
# initial/boundary conditions
v = np.zeros((nr+1,nz+1))
v[:,0] = 1100
v[:,-1] = 0
v[31:,29:40] = 1000
v[19:,54:65] = -200
# convergence parameters
tol = 1e-4
# work arrays
v_old = np.empty_like(v)
w1 = np.empty_like(v[0, 1:nz])
w2 = np.empty_like(v[1:nr,1:nz])
w3 = np.empty_like(v[nr, 1:nz])
# constants
A = cr * (1 - 1/(2*ri[1:nr,1:nz]))
B = cr * (1 + 1/(2*ri[1:nr,1:nz]))
# Gauss-Seidel solver
tic = time.time()
max_v_diff = 1;
while (max_v_diff > tol):
v_old[:] = v
# left boundary updates
np.add(v_old[0, 0:nz-1], v_old[0, 2:nz+2], out=v[0, 1:nz])
v[0, 1:nz] *= cz
np.multiply(2*cr, v_old[1, 1:nz], out=w1)
v[0, 1:nz] += w1
# internal updates
np.add(v_old[1:nr, 0:nz-1], v_old[1:nr, 2:nz+1], out=v[1:nr, 1:nz])
v[1:nr,1:nz] *= cz
np.multiply(A, v_old[0:nr-1, 1:nz], out=w2)
v[1:nr,1:nz] += w2
np.multiply(B, v_old[2:nr+1, 1:nz], out=w2)
v[1:nr,1:nz] += w2
# right boundary updates
np.add(v_old[nr, 0:nz-1], v_old[nr, 2:nz+1], out=v[nr, 1:nz])
v[nr, 1:nz] *= cz
np.multiply(2*cr, v_old[nr-1, 1:nz], out=w3)
v[nr,1:nz] += w3
# reapply grid potentials
v[31:,29:40] = 1000
v[19:,54:65] = -200
# check for convergence
v_old -= v
max_v_diff = np.absolute(v_old).max()
toc = time.time() - tic
- 优化“内部更新”行。
“内部更新”行是求解器中最耗时的部分。可以使用NumPy的编译工具之一scipy.weave.blitz来编译NumPy表达式,如“内部更新”行,以获得更快的速度。以下是优化的代码:
from scipy import weave
# [...] Set up code till "# Gauss-Seidel solver"
tic = time.time()
max_v_diff = 1;
A = cr * (1 - 1/(2*ri[1:nr,1:nz]))
B = cr * (1 + 1/(2*ri[1:nr,1:nz]))
expr = "temp = A*v[0:nr-1,1:nz] + B*v[2:nr+1,1:nz] + cz*(v[1:nr,0:nz-1] + v[1:nr,2:nz+1]); v[1:nr,1:nz] = temp"
temp = np.empty((nr-1, nz-1))
while (max_v_diff > tol):
v_old = v.copy()
# left boundary updates
v[0,1:nz] = cr*2*v[1,1:nz] + cz*(v[0,0:nz-1] + v[0,2:nz+2])
# internal updates
weave.blitz(expr, check_size=0)
# right boundary updates
v[nr,1:nz] = cr*2*v[nr-1,1:nz] + cz*(v[nr,0:nz-1] + v[nr,2:nz+1])
# reapply grid potentials
v[31:,29:40] = 1000
v[19:,54:65] = -200
# check for convergence
v_diff = v - v_old
max_v_diff = np.absolute(v_diff).max()
toc = time.time() - tic
经过这些优化后,NumPy版本的求解器速度得到了显著提升,与MATLAB版本基本相当。