如何用numpy向量化替代域循环，排查残差不变问题？

我需要遍历(i,j)域进行数学运算，但嵌套循环运行耗时极长，因此希望用向量化计算替代for循环。已尝试用numpy实现向量化改写，但完整版本中残差（residual）完全没有变化，怀疑sigma未进行正确迭代，不确定问题出在哪里。

原循环代码 #

import numpy as np

# Variable Init
tole = 1e-4
residual = 1
Nx = int(1e3)
Ny = Nx
sigma = np.zeros((Nx,Ny))
R = np.zeros((Nx,Ny))

# Sample constants
a = 10
b = 3
c = 2
d = 1
w = 0.5

# Output: Sigma and residual after convergence
while residual > tole:
    for i in range(1,Nx-1):
        for j in range(1,Ny-1):
            temp = a * sigma[j, i + 1] + b * sigma[j, i - 1] + c * sigma[j + 1, i] + d * sigma[j - 1, i]

            R[j, i] = abs(sigma[j,i] - temp)
            sigma[j, i] = (1 - w) * sigma[j,i] + w * temp
    residual = np.sum(R)/(Nx*Ny)

改写后的错误numpy代码 #

import numpy as np

# Variable Init
tole = 1e-4
residual = 1
Nx = int(1e3)
Ny = Nx
sigma = np.zeros((Nx,Ny))
R = np.zeros((Nx,Ny))

# Sample constants
a = 10
b = 3
c = 2
d = 1
w = 0.5

# Output: Sigma and residual after convergence
while residual > tole:
    
    temp = a * sigma[:-1, 1:] + b * sigma[1:, :-1] + c * sigma[1:, :-1] + d * sigma[:-1, 1:]
    R = sigma[:-1,:-1] - temp
    sigma[:-1,:-1] = (1 - w) * sigma[:-1,:-1] + w * temp
    residual = np.sum(R)/(Nx*Ny)

问题分析与修正 #

你的向量化代码存在三个核心错误，导致sigma迭代逻辑完全偏离原循环：

• 邻域切片完全错误：原循环中temp是当前点的右、左、下、上四个邻域的加权和，但你写的切片没有对应正确的邻域，甚至重复使用了错误的切片（c和d的项和b、a的项重复）。正确的切片对应关系是：
  • sigma[j, i+1]（右侧）→ sigma[1:-1, 2:]
  • sigma[j, i-1]（左侧）→ sigma[1:-1, :-2]
  • sigma[j+1, i]（下侧）→ sigma[2:, 1:-1]
  • sigma[j-1, i]（上侧）→ sigma[:-2, 1:-1]
• 未计算残差的绝对值：原循环中R[j,i]取了abs(sigma[j,i]-temp)，但你改写时漏掉了绝对值，导致残差计算错误。
• sigma更新区域错误：原循环只更新i=1~Nx-2、j=1~Ny-2的中间区域，对应向量化的sigma[1:-1,1:-1]，而你用的sigma[:-1,:-1]会包含第0行和第0列的边界点，不符合原逻辑。

修正后的向量化代码 #

import numpy as np

# Variable Init
tole = 1e-4
residual = 1
Nx = int(1e3)
Ny = Nx
sigma = np.zeros((Nx,Ny))
R = np.zeros((Nx,Ny))

# Sample constants
a = 10
b = 3
c = 2
d = 1
w = 0.5

# Output: Sigma and residual after convergence
while residual > tole:
    # 取出中间需要更新的区域
    sigma_mid = sigma[1:-1, 1:-1]
    # 计算四个邻域的加权和
    temp = a * sigma[1:-1, 2:] + b * sigma[1:-1, :-2] + c * sigma[2:, 1:-1] + d * sigma[:-2, 1:-1]
    # 计算残差（取绝对值）
    R[1:-1, 1:-1] = np.abs(sigma_mid - temp)
    # 更新中间区域的sigma
    sigma[1:-1, 1:-1] = (1 - w) * sigma_mid + w * temp
    # 计算平均残差
    residual = np.sum(R) / (Nx * Ny)

内容的提问来源于stack exchange，提问作者Winston Zhang