3

这个问题是最近发布的关于 MATLAB 的速度是 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
4

2 回答 2

2

在我的笔记本电脑上，您的代码在大约 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
于 2013-07-10T21:53:24.253 回答
2

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

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

  2. 选择一种方法来优化它。有几种工具可以加速 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)

  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
于 2013-07-10T23:35:11.677 回答