from mpi4py import MPI

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


# unicast

if rank == 0:
    data = [1, 2, 3, 4]
    comm.send(data, dest=1, tag=7)
elif rank == 1:
    data = comm.recv(source=0, tag=7)
    print(‘Received', data)


# broadcast

if rank == 0:
    data = {'a':1,'b':2,'c':3}
else:
    data = None

data = comm.bcast(data, root=0)
print(data)


# scatter

if rank == 0:
    data = np.arange(20)
    chunks = np.array_split(data, size)
else:
    chunks = None

chunk = comm.scatter(chunks, root=0)

# size = 5,  rank = 2:
# chunk = [8, 9, 10, 11]
# rank = 0:
# chunk = [0, 1, 2, 3]


# gather

chunk = comm.scatter(chunks, root=0) 

results = []
for x in chunk:
    results.append(x*x)

gathered_chunks = comm.gather(results, root=0)

# rank = 0:
# g_c = [ [0, 1, 4, 9], [16, 25, 36, 49], … ]
# rank != 0:
# g_c = None