CoreFlows/examples/Python/MPI4PY/testMPI4PY.py

   1 #!/usr/bin/env python3
   2 # -*-coding:utf-8 -*
   3
   4 #===============================================================================================================================
   5 # Name        : Tests of the library mpi4py from MPI4PY tutorial
   6 # Author      : Michaël Ndjinga
   7 # Copyright   : CEA Saclay 2021
   8 # Description : https://mpi4py.readthedocs.io/en/stable/tutorial.html
   9 #================================================================================================================================
  10
  11 from mpi4py import MPI
  12 import numpy as np
  13
  14 # Tests from MPI4PY tutorial https://mpi4py.readthedocs.io/en/stable/tutorial.html
  15
  16 comm = MPI.COMM_WORLD
  17 size = comm.Get_size()
  18 rank = comm.Get_rank()
  19
  20 print("My rank is ", rank, " among ", size, "processors ")
  21
  22 ###Point-to-Point Communication
  23
  24 #Python objects (pickle under the hood):
  25
  26 if rank == 0:
  27     data = {'a': 7, 'b': 3.14}
  28     comm.send(data, dest=1, tag=11)
  29 elif rank == 1:
  30     data = comm.recv(source=0, tag=11)
  31
  32 #Python objects with non-blocking communication:
  33
  34 if rank == 0:
  35     data = {'a': 7, 'b': 3.14}
  36     req = comm.isend(data, dest=1, tag=11)
  37     req.wait()
  38 elif rank == 1:
  39     req = comm.irecv(source=0, tag=11)
  40     data = req.wait()
  41
  42 # passing MPI datatypes explicitly
  43 if rank == 0:
  44     data = np.arange(1000, dtype='i')
  45     comm.Send([data, MPI.INT], dest=1, tag=77)
  46 elif rank == 1:
  47     data = np.empty(1000, dtype='i')
  48     comm.Recv([data, MPI.INT], source=0, tag=77)
  49
  50 # automatic MPI datatype discovery
  51 if rank == 0:
  52     data = np.arange(100, dtype=np.float64)
  53     comm.Send(data, dest=1, tag=13)
  54 elif rank == 1:
  55     data = np.empty(100, dtype=np.float64)
  56     comm.Recv(data, source=0, tag=13)
  57
  58 ###Collective Communication
  59
  60 #Broadcasting a Python dictionary:
  61
  62 if rank == 0:
  63     data = {'key1' : [7, 2.72, 2+3j],
  64             'key2' : ( 'abc', 'xyz')}
  65 else:
  66     data = None
  67 data = comm.bcast(data, root=0)
  68
  69 #Scattering Python objects:
  70
  71 if rank == 0:
  72     data = [(i+1)**2 for i in range(size)]
  73 else:
  74     data = None
  75 data = comm.scatter(data, root=0)
  76 assert data == (rank+1)**2
  77
  78 #Gathering Python objects:
  79
  80 data = (rank+1)**2
  81 data = comm.gather(data, root=0)
  82 if rank == 0:
  83     for i in range(size):
  84         assert data[i] == (i+1)**2
  85 else:
  86     assert data is None
  87
  88 # Broadcasting a NumPy array:
  89
  90 if rank == 0:
  91     data = np.arange(100, dtype='i')
  92 else:
  93     data = np.empty(100, dtype='i')
  94 comm.Bcast(data, root=0)
  95 for i in range(100):
  96     assert data[i] == i
  97
  98 #Scattering NumPy arrays:
  99
 100 sendbuf = None
 101 if rank == 0:
 102     sendbuf = np.empty([size, 100], dtype='i')
 103     sendbuf.T[:,:] = range(size)
 104 recvbuf = np.empty(100, dtype='i')
 105 comm.Scatter(sendbuf, recvbuf, root=0)
 106 assert np.allclose(recvbuf, rank)
 107
 108 #Gathering NumPy arrays:
 109
 110 sendbuf = np.zeros(100, dtype='i') + rank
 111 recvbuf = None
 112 if rank == 0:
 113     recvbuf = np.empty([size, 100], dtype='i')
 114 comm.Gather(sendbuf, recvbuf, root=0)
 115 if rank == 0:
 116     for i in range(size):
 117         assert np.allclose(recvbuf[i,:], i)
 118
 119 #Parallel matrix-vector product:
 120
 121