MPI Gather Function in Python

The gather function is used to gather data from multiple processes into a single process. We’ll go through the provided code, line by line, and understand how the gather function works. The detailed tutorial can be found here.

Code

from mpi4py import MPI

comm = MPI.COMM_WORLD
size = comm.Get_size()
rank = comm.Get_rank()
data = (rank + 1)
data = comm.gather(data, root=0)
print(f"Process {rank}: Calculated data = {data}")
if rank == 0:
    for i in range(size):
        assert data[i] == (i + 1)
        print(f"Process 0: Checked data received from process {i + 1}")
else:
    assert data is None

Explanation

from mpi4py import MPI

This line imports the MPI functionality from the mpi4py library.

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

These lines initialize the MPI communicator (comm) and obtain the total number of processes (size) and the rank of the current process (rank).

data = (rank + 1)
data = comm.gather(data, root=0)

Each process calculates its own data value based on its rank (rank + 1). Then, the gather function is called on the communicator comm. This function gathers data from all processes and…