掌握mpi4py中不同形状数组的聚合操作

from mpi4py import MPI
import numpy as np

comm = MPI.COMM_WORLD
size = comm.Get_size()
rank = comm.Get_rank()

# 模拟不同形状的数组
a = np.zeros((2 if rank == 1 else 5, 3), dtype=float) + rank
print(f"Rank {rank}: 数组形状 {a.shape}")

# comm.Gather 会在这里失败，因为它期望所有进程发送相同大小的数据
# b = np.zeros((12, 3), dtype=float) - 1 # 预分配一个足够大的缓冲区，但comm.Gather仍会因形状不匹配而失败
# comm.Gather(a, b, root=0)
# if rank == 0:
#     print(b)

import numpy as np
from mpi4py import MPI

comm = MPI.COMM_WORLD
size = comm.Get_size()
rank = comm.Get_rank()

# 模拟不同形状的数组
a = np.zeros((2 if rank == 1 else 5, 3), dtype=float) + rank
print(f"Rank {rank}: 数组形状 {a.shape}, 数组内容:\n{a}")

# 使用 comm.gather 聚合数组
# 根进程会收到一个包含所有进程发送数组的列表
gathered_arrays = comm.gather(a, root=0)

if rank == 0:
    # 根进程将收到的数组列表进行拼接
    concatenated_array = np.concatenate(gathered_arrays)
    print(f"\nRank {rank}: 聚合并拼接后的数组形状 {concatenated_array.shape}, 数组内容:\n{concatenated_array}")
else:
    print(f"\nRank {rank}: 非根进程不接收聚合结果")

import numpy as np
from mpi4py import MPI

comm = MPI.COMM_WORLD
size = comm.Get_size()
rank = comm.Get_rank()

# 确保进程数量不超过2，以便示例清晰
assert size <= 2, "此Gatherv示例仅为2个进程设计"

if rank == 0:
    a = np.zeros((5, 3), dtype=float) + rank
else:
    a = np.zeros((2, 3), dtype=float) + rank

print(f"Rank {rank}: 数组形状 {a.shape}, 数组内容:\n{a}")

# 计算全局总行数
# 对于 rank 0 (5行) 和 rank 1 (2行)，总共 7 行
n_global_rows = 7 
n_cols = a.shape[1] # 列数保持不变

# 根进程预分配接收缓冲区
if rank == 0:
    b = np.zeros((n_global_rows, n_cols), dtype=float)

    # 定义每个进程发送的元素数量 (行数 * 列数)
    # 假设 rank 0 发送 5*3=15 个元素，rank 1 发送 2*3=6 个元素
    send_counts = [5 * n_cols, 2 * n_cols] 

    # 定义每个进程数据在 b 中的起始偏移量 (以元素为单位)
    # rank 0 的数据从 b[0] 开始 (0个元素偏移)
    # rank 1 的数据从 b[5*n_cols] 开始 (15个元素偏移，即在 rank 0 数据之后)
    displacements = [0, 5 * n_cols] 

    # 接收缓冲区描述符
    recvbuf_params = (b, send_counts, displacements, MPI.DOUBLE)
else:
    b = None # 非根进程不需要接收缓冲区，设为None
    recvbuf_params = None # 非根进程也不需要接收缓冲区参数

# 执行 Gatherv 操作
comm.Gatherv(a, recvbuf_params, root=0)

if rank == 0:
    print(f"\nRank {rank}: Gatherv 聚合后的数组形状 {b.shape}, 数组内容:\n{b}")
else:
    print(f"\nRank {rank}: 非根进程不接收 Gatherv 结果")