提升 NumPy 高斯-赛德尔（雅可比）求解器的速度

一位用户在将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版本基本相当。