如何在Python中使用多进程遍历x*y网格域并执行计算？

我需要编写代码遍历x*y网格点并执行计算，示例代码如下：

Nx = 10
Ny = 10
for i in range(len(1,Nx-1))
    for j in range(len(1, Ny-1))
        temp = a*sigma[i,j]+b*sigma[i+1,j]+c*sigma[i,j+1]
        R = sigma[i,j]-temp

示例用10做演示值，但实际域规模大得多，遍历计算耗时久。尝试在循环中嵌套多进程时，因需要索引j+1等值遇到问题，我的多进程代码如下：

import numpy as np
from multiprocessing import Pool

def mathstuff(a,b,c,sigma,j):
    temp = a*sigma[i,j]+b*sigma[i+1,j]+c*sigma[i,j+1] + c*sigma[i,j-1]
    R = sigma[i,j]-temp

Nx = 10
Ny = 10
for i in range(len(1,Nx-1))
    with Pool() as pool:
        j=np.arange(1,Ny-1,1)
        results = mathstuff(a,b,c,sigma,j),j)

报错提示numpy数组不可索引，不清楚该如何解决。

核心错误点 #

• range语法错误：range(len(1,Nx-1))写法错误，range只需传入起始和终止值，应为range(1, Nx-1)。
• 变量作用域问题：mathstuff函数引用了外部的i变量，多进程中子进程无法直接访问该变量，会导致未定义错误。
• 索引逻辑错误：将整个j数组传入函数后，sigma[i,j]会尝试用数组索引数组，不符合numpy的索引规则，引发报错。
• 多进程调用语法错误：results = mathstuff(a,b,c,sigma,j),j)存在语法错误，正确的多进程调用需使用pool.map或类似方法。

修正后的多进程实现 #

推荐按行拆分任务（减少进程间数据传递开销），代码如下：

import numpy as np
from multiprocessing import Pool

def process_row(args):
    i, a, b, c, sigma, Ny = args
    row_R = np.zeros(Ny)
    # 遍历当前行的有效j范围
    for j in range(1, Ny-1):
        temp = a * sigma[i,j] + b * sigma[i+1,j] + c * sigma[i,j+1] + c * sigma[i,j-1]
        row_R[j] = sigma[i,j] - temp
    return i, row_R

if __name__ == "__main__":
    # 初始化参数
    a, b, c = 1.0, 0.5, 0.3
    Nx, Ny = 100, 100  # 实际大尺寸示例
    sigma = np.random.rand(Nx, Ny)  # 模拟输入数组
    R = np.zeros_like(sigma)

    # 构造任务列表：传入每一行的i和必要参数
    tasks = [(i, a, b, c, sigma, Ny) for i in range(1, Nx-1)]
    
    # 启动多进程池
    with Pool() as pool:
        # 并行处理每一行
        results = pool.map(process_row, tasks)
    
    # 将结果填充回R数组
    for i, row_R in results:
        R[i] = row_R

更高效的替代方案：numpy向量化运算 #

numpy的向量化操作底层经过优化，速度通常比多进程循环更快，尤其适合大规模数组：

import numpy as np

# 初始化参数
a, b, c = 1.0, 0.5, 0.3
Nx, Ny = 100, 100
sigma = np.random.rand(Nx, Ny)
R = np.zeros_like(sigma)

# 向量化计算，直接用数组切片
R[1:Nx-1, 1:Ny-1] = sigma[1:Nx-1, 1:Ny-1] - (
    a * sigma[1:Nx-1, 1:Ny-1] +
    b * sigma[2:Nx, 1:Ny-1] +
    c * sigma[1:Nx-1, 2:Ny] +
    c * sigma[1:Nx-1, 0:Ny-2]
)

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