Python 提高Gauss-Seidel（Jacobi）求解器的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秒）

我尝试了从中的解决方案，将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

---

在我的笔记本电脑上，你的代码运行大约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

---

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

找到瓶颈。我从控制台中分析了代码，以查找花费时间最多的行。事实证明，80%以上的时间都花在了“内部更新”这一行上

选择一种方法来优化它。在numpy中有几种工具可以加速代码（Cython、numexpr、weave…）。特别是，

scipy.weave.blitz

非常适合将numpy表达式（如有问题的行）编译成快速代码。理论上，该行可以被包装在

“…”

中，并作为

weave.blitz（“…”）

执行，但要在计算中使用正在更新的数组，因此，如图中第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)

检查结果是否正确后，使用

weave.blitz（expr，check\u size=0）

禁用运行时检查。代码现在在34秒内运行

以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

---

您可以从分析代码和确定瓶颈开始。如果您希望提高性能，您是否考虑过Cython（编译成C的类似Python的代码）或Numba（使用LLVM的JIT编译）？这里有一个有趣的比较：我应该提到，这个实现实际上是Jacobi而不是Gauss-Seidel，并且存在应用潜力（但这些潜力可以很容易地消除）。这是一个非常好的方式来展示临时数组的昂贵程度。这可能会导致代码的可读性降低，但我会尽量记住使用

out=…

参数的可能性。是的，这表明了同样的问题（jaime也给出了一个有趣的回答；-），即一旦开始在临时数组上使用大量ram，计算时间实际上会快速增长，通过使用A和B的技巧，我的机器上的计算时间从55秒增加到28秒。（请注意，在Matlab上执行同样的操作也会减少计算时间。）在我的机器上，所有其他使用临时数组的操作都比较慢。好的，因此A和B方法提供了巨大的改进：Python中的30s->16.2s和Matlab中的20s->13s。这让Python迎头赶上！但是，Jaime，您使用临时数组提供的代码实际上比前面的评论员建议的要慢。使用此方法使Python在21.8s中进行模拟。这有点令人费解。。。不反复计算

A

和

B

有一个明显的优势，但在numpy中有一些不太正确的地方。试着计时下面的内容来了解我的意思：

a=np.arange（1000）；b=np.arange（1000）；c=np.空的_像（a）；%timeit a+b（100000个循环，每个循环的最佳值为3:1.97 us）；%timeit np.添加（a，b）（1000000个循环，3个循环的最佳值：每个循环2.94 us）；timeit np.add（a，b，out=c）（10000个循环，每个循环最好3:2.22 us）

a+b

应该是

np的简写。加上（a，b）

，很难理解开销从何而来……反过来，使用常数我得到28秒而不是55秒，使用weave我得到24秒。主要的改进实际上是通过使用常量@这很有趣：我还没有测试过这种组合，而且速度确实很快。不过，在我的笔记本电脑中，最终的weave版本比你的计时快了30%，而不是15%。顺便问一下，你能检查一下Matlab在使用A和B时是否真的更快吗？也许它能认识到它们没有变化？我在Matlab中实现了A和B的变化，时间从20秒到13秒，这是一个明显的改进。谢谢你的代码。现在，约40%的时间用于运行

weave.blitz

，约30%的时间用于检查收敛性，因此，加快收敛速度的一个简单方法（当您需要多次迭代时）是仅每隔一次迭代检查收敛性。