在使用mpi4py进行并行计算时,我们通常需要从所有子进程那里收集结果。以下是一个简单的Python代码示例,展示了如何使用mpi4py收集子进程的结果:
from mpi4py import MPI
import numpy as np
comm = MPI.COMM_WORLD
size = comm.Get_size()
rank = comm.Get_rank()
# 假设每个子进程都计算一个简单的函数值
def calculate_value(rank):
return rank * np.pi
# 主进程等待所有子进程计算完毕,并收集结果
if rank == 0:
results = np.empty(size, dtype=float)
for i in range(size):
results[i] = comm.recv(source=i, tag=1)
print(results)
else:
comm.send(calculate_value(rank), dest=0, tag=1)
在这个例子中,主进程(rank为0的进程)负责收集其他子进程的结果。每个子进程在计算完毕后,使用comm.send发送结果给主进程,主进程通过循环使用comm.recv接收来自各个子进程的结果,并将它们存储在results数组中。