提高 Gauss-Seidel (Jacobi) 求解器的 Numpy 速度

Question

此问题是recent question posted regarding MATLAB being twice as fast as Numpy 的后续问题。

我目前在 MATLAB 和 Numpy 中实现了一个 Gauss-Seidel 求解器，它作用于二维轴对称域（圆柱坐标）。代码最初是用 MATLAB 编写的，然后转移到 Python 中。 Matlab 代码运行时间约为 20 秒，而 Numpy 代码运行时间约为 30 秒。但是，我想使用 Numpy，因为这段代码是一个更大程序的一部分，几乎两倍长的模拟时间是一个很大的缺点。

该算法只是求解矩形网格上的离散拉普拉斯方程（在圆柱坐标中）。当网格更新之间的最大差异小于指定的容差时，它完成。

Numpy 中的代码是：

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

# Gauss-Seidel solver
tic = time.time()
max_v_diff = 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
    v[1:nr,1:nz] = cr*((1 - 1/(2*ri[1:nr,1:nz]))*v[0:nr-1,1:nz] + (1 + 1/(2*ri[1:nr,1:nz]))*v[2:nr+1,1:nz]) + cz*(v[1:nr,0:nz-1] + v[1:nr,2:nz+1])
    # 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
print(toc)

这实际上不是我使用的完整算法。完整的算法使用连续的过松弛和棋盘迭代方案来提高速度并消除求解器的方向性，但为了简单起见，我提供了这个更易于理解的版本。 Numpy 的速度缺陷在完整版中更为明显（Numpy 和 MATLAB 中的仿真时间分别为 17 秒和 9 秒）。

我尝试了previous question 中的解决方案，将 v 更改为以列为主的顺序数组，但没有性能提升。

有什么建议吗？

编辑：Matlab 代码供参考：

% geometry
length = 0.008;
width = 0.002;

% mesh
nz = 256;
nr = 64;

% step sizes
dz = length/nz;
dr = width/nr;

% node position matrices
r = repmat(linspace(0,width,nr+1)', 1, nz+1);
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 = zeros(nr+1,nz+1);
v(1:nr+1,1) = 1100;
v(1:nr+1,nz+1) = 0;
v(32:nr+1,30:40) = 1000;
v(20:nr+1,55:65) = -200;

% convergence parameters
tol = 1e-4;
max_v_diff = 1;

% Gauss-Seidel Solver
tic
while (max_v_diff > tol)
    v_old = v;

    % left boundary updates
    v(1,2:nz) = cr.*2.*v(2,2:nz) + cz.*( v(1,1:nz-1) + v(1,3:nz+1) );
    % internal updates
    v(2:nr,2:nz) = cr.*( (1 - 1./(2.*ri(2:nr,2:nz))).*v(1:nr-1,2:nz) + (1 + 1./(2.*ri(2:nr,2:nz))).*v(3:nr+1,2:nz) ) + cz.*( v(2:nr,1:nz-1) + v(2:nr,3:nz+1) );    
    % right boundary updates
    v(nr+1,2:nz) = cr.*2.*v(nr,2:nz) + cz.*( v(nr+1,1:nz-1) + v(nr+1,3:nz+1) );
    % reapply grid potentials
    v(32:nr+1,30:40) = 1000;
    v(20:nr+1,55:65) = -200;

    % check for convergence
    max_v_diff = max(max(abs(v - v_old)));

end
toc

Answer 1

在我的笔记本电脑上，您的代码运行大约需要 45 秒。通过尝试将中间数组的创建减少到最低限度，包括重用预先分配的工作数组，我设法将时间减少到 27 秒。这应该会让你回到 MATLAB 的水平，但你的代码可读性会降低。不管怎样，找到下面的代码来替换你# Gauss-Seidel solver评论下面的所有内容：

# 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

Answer 2

通过遵循以下流程，我已经能够将笔记本电脑的运行时间从 66 秒减少到 21 秒：

1. 找到瓶颈。我使用IPython 控制台中的line_profiler 对代码进行了概要分析，以找到花费最多时间的行。事实证明，超过 80% 的时间都花在了“内部更新”这一行上。

2. 选择一种优化方法。有几种工具可以加速 numpy 中的代码（Cython、numexpr、weave...）。特别是，scipy.weave.blitz 非常适合将 numpy 表达式（如违规行）编译为快速代码。理论上，该行可以包含在"..." 中并作为weave.blitz("...") 执行，但是正在更新的数组用于计算，因此如docs 中的第4 点所述，必须使用一个临时数组来保持同样的结果：

   expr = "temp = cr*((1 - 1/(2*ri[1:nr,1:nz]))*v[0:nr-1,1:nz] + (1 + 1/(2*ri[1:nr,1:nz]))*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 ...
       # internal updates
       weave.blitz(expr)
   

3. 检查结果是否正确后，使用weave.blitz(expr, check_size=0) 禁用运行时检查。代码现在在 34 秒内运行。

4. 在 Jaime 的工作基础上，预先计算表达式中的常数因子 A 和 B。代码在 21 秒内运行（改动很小，但现在需要编译器）。

这是代码的核心：

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