如何使用Multiprocessing在Python中循环遍历x*y域？

Question

我正在尝试编写一些循环遍历 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 数组将不可索引，这是公平的，但我似乎不知道如何解决这个问题。

Answer 1

下面是包含如何使用多处理的解决方案的代码。代码中需要注意以下几点：

1. 如果

   sigma

   是一个非常大的数组，请考虑将其作为共享数据传递（请注意，我正在谈论使用数组作为只读）。否则，您将在每个线程中创建

   sigma

   数组的副本。
2. 参数在多处理中作为列表/元组传递给函数，因此您需要在函数中将它们解压

   a,b,c,i,j = args

   。
3. 对于多处理实现，您仍然需要创建 (i,j) 的所有组合并将它们作为参数传递给池。
4. 要将多处理应用于函数，您需要将函数映射到池，其中包含要应用于函数的所有参数的列表： `pool.map(your_func, [(arg1,arg2,...),(arg1, arg2,...),...])

import numpy as np
from multiprocessing import Pool
from itertools import product


# initialize shared memory for worker processes
def func_init_shared(data: np.ndarray):
    # declare global variable 
    # that will be shared across all workers in multiprocessing pool
    global shared_data

    # store our numpy array in the global variable
    shared_data = data


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


if __name__ == '__main__':
        
    Nx = 10
    Ny = 10
    a,b,c = -1, -2, -3

    np.random.seed(3)
    simga = np.random.rand(10,10)

    indexes = list(product(range(1,Nx-1), range(1,Ny-1)))

    with Pool(initializer=func_init_shared, initargs=(simga,)) as pool:
        # each task is a set of arguments to your function
        tasks = [(a,b,c,idx[0],idx[1]) for idx in indexes]

        results = pool.map(mathstuff, tasks)   
        print(results)