mpi hello world

hello_mpi.py

"""

Running this example:
    mpirun -np 4 python MPI_in_python.py
    mpirun -np 4 --hostfile MPI_hostfile python MPI_in_python.py

"""

import logging
import numpy as np
from mpi4py import MPI

process_host = MPI.Get_processor_name()
comm = MPI.COMM_WORLD

num = (100, 1000)

logging.basicConfig(
    format=(
        "{asctime} - rank"
        f"{comm.rank}"
        f" @ {process_host}"
        " - {name} - {levelname} - {message}"
    ),
    datefmt="%Y-%m-%dT%H:%M:%S%z",
    style="{",
    level=logging.INFO,
)
logger = logging.getLogger("mpi-example")

rank_inds = [np.arange(rank, num[0], comm.size) for rank in range(comm.size)]
inds = rank_inds[comm.rank]
if comm.rank == 0:
    logger.info("generating data")
    np.random.seed(8374)
    data = np.abs(np.random.randn(*num))

    for rank in range(1, comm.size):
        logger.info(f"sending data to rank {rank}")

        comm.send(data[rank_inds[rank], :], dest=rank, tag=0)

    subdata = data[inds, :]
    del data
else:
    logger.info("receiving data from rank 0")
    subdata = comm.recv(source=0, tag=0)

logger.info("computing sum...")
sub_results = np.empty(len(inds))
for ind in range(len(inds)):
    sub_results[ind] = np.sum(subdata[ind, :])
logger.info("computing done")

if comm.rank == 0:
    results = np.empty(num[0])

    logger.info("collecting results in rank 0")

    results[inds] = sub_results
    for rank in range(1, comm.size):
        results[rank_inds[rank]] = comm.recv(source=rank, tag=rank)
        logger.info(f"received data from rank {rank}")

    print(results)
else:
    logger.info(f"sending data from rank {comm.rank}")
    comm.send(sub_results, dest=0, tag=comm.rank)