Newer
Older
from __future__ import print_function
from mpi4py import MPI
import numpy as np
comm = MPI.COMM_WORLD
size = comm.Get_size()
rank = comm.Get_rank()
if size != 2:
raise RuntimeError("Please run with two MPI tasks")
data = np.zeros((8,8), int)
if rank == 0:
for i in range(8):
data[i,:] = np.arange(1, 9) + (i+1) * 10
if rank == 0:
print("Original data:")
print(data)
# TODO Create the custom datatype
# Note: Python integers are 64 bits
# TODO: communicate with the datatype
# Note: mpi4py requires that the input and output buffers are contiguous
# in memory. Thus, in order to send/receive from second column we need
# to create a contiguous view starting from the second element in memory
# which can be accomplished by flattening the array into 1d with ravel
if rank == 1:
print("Received data:")
print(data)