python - 提高 Gauss-Seidel (Jacobi) Solver 的 Numpy 速度

Question

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

我目前有一个在 MATLAB 和 Numpy 中实现的高斯-赛德尔求解器，它作用于二维轴对称域(圆柱坐标)。代码最初是用 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

按照以下过程，我已经能够将笔记本电脑的运行时间从 66 秒减少到 21 秒:

1. 找到瓶颈。我使用 line_profiler 分析了代码来自 IPython控制台查找花费最多时间的行。事实证明，超过 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