Clean up flushing, add communication of Python objects

Author: jussienko

demos/mpi-alltoall.py

@@ -1,6 +1,6 @@
 from mpi4py import MPI
 import numpy as np
-import sys
+from sys import stdout
 
 comm = MPI.COMM_WORLD
 rank = comm.Get_rank()
@@ -10,6 +10,10 @@
 
 data = np.arange(8) / 10. + rank
 recv_buf = np.zeros(8)
+# Python sequence, lenght has to be equal number to MPI tasks
+py_data = []
+for r in range(4):
+    py_data.append({'key{0:02d}'.format(10*rank + r) : 10*rank + r})
 
 if rank == 0:
     print("Original data")
@@ -18,22 +22,28 @@
 for r in range(size):
     if rank == r:
         print("rank ", rank, data)
-        sys.stdout.flush()
+        print("rank ", rank, py_data)
     comm.Barrier()
 
+stdout.flush()
+
 comm.Alltoall(data, recv_buf)
+new_data = comm.alltoall(py_data)
 
 comm.Barrier()
 if rank == 0:
     print()
     print("Final data")
-    sys.stdout.flush()
+
+stdout.flush()
 comm.Barrier()
 
 for r in range(size):
     if rank == r:
         print("rank ", rank, recv_buf)
-        sys.stdout.flush()
+        print("rank ", rank, new_data)
     comm.Barrier()
 
+stdout.flush()
+
 

demos/mpi-bcast.py

@@ -1,6 +1,6 @@
 from mpi4py import MPI
 import numpy as np
-import sys
+from sys import stdout
 
 comm = MPI.COMM_WORLD
 rank = comm.Get_rank()
@@ -9,27 +9,32 @@
 assert size == 4
 
 if rank == 0:
-    data = np.arange(8) / 10.
+    data = np.arange(8) / 10. # NumPy array
+    py_data = {'key1' : 0.0, 'key2' : 11} # Python object
 else:
     data = np.zeros(8)
+    py_data = None
 
 if rank == 0:
     print("Original data")
-    sys.stdout.flush()
+stdout.flush()
 comm.Barrier()
 
 print("rank ", rank, data)
-sys.stdout.flush()
+print("rank ", rank, py_data)
+stdout.flush()
 
 comm.Bcast(data, root=0)
+new_data = comm.bcast(py_data, root=0)
 
 comm.Barrier()
 if rank == 0:
     print()
     print("Final data")
-    sys.stdout.flush()
+stdout.flush()
 comm.Barrier()
 
 print("rank ", rank, data)
+print("rank ", rank, new_data)
 
 

demos/mpi-gather.py

@@ -1,37 +1,52 @@
 from mpi4py import MPI
 import numpy as np
-import sys
+from sys import stdout
 
 comm = MPI.COMM_WORLD
 rank = comm.Get_rank()
 size = comm.Get_size()
 
 assert size == 4
 
-data = np.arange(2) / 10. + rank
+data = np.arange(2) / 10. + rank  # NumPy array
 
+# Let's create different Python objects for different MPI tasks
 if rank == 0:
+    py_data = 'foo.bar'
+elif rank == 1:
+    py_data = 12.34
+elif rank == 2:
+    py_data = {'key1' : 99.0, 'key2' : [-1, 2.3]}
+else:
+    py_data = [6.5, 4.3]
+
+
+if rank == 1:
     recv_buf = np.zeros(8)
 else:
     recv_buf = None
 
 if rank == 0:
     print("Original data")
-    sys.stdout.flush()
+
+stdout.flush()
 comm.Barrier()
 
 print("rank ", rank, data)
-sys.stdout.flush()
+print("rank ", rank, py_data)
+stdout.flush()
 
-comm.Gather(data, recv_buf, root=0)
+comm.Gather(data, recv_buf, root=1)
+new_data = comm.gather(py_data, root=1)
 
 comm.Barrier()
 if rank == 0:
     print()
     print("Final data")
-    sys.stdout.flush()
+
+stdout.flush()
 comm.Barrier()
 
 print("rank ", rank, recv_buf)
-
+print("rank ", rank, new_data)
 

demos/mpi-scatter.py

@@ -1,6 +1,6 @@
 from mpi4py import MPI
 import numpy as np
-import sys
+from sys import stdout
 
 comm = MPI.COMM_WORLD
 rank = comm.Get_rank()
@@ -9,29 +9,37 @@
 assert size == 4
 
 if rank == 0:
-    data = np.arange(8) / 10.
+    data = np.arange(8) / 10.  # NumPy array
+    # Python sequence, lenght has to be equal number to MPI tasks
+    py_data = ['foo', 'bar', 11.2, {'key' : 22}]
 else:
     data = None
+    py_data = None
 
 recv_buf = np.zeros(2)
 
 if rank == 0:
     print("Original data")
-    sys.stdout.flush()
+
+stdout.flush()
 comm.Barrier()
 
 print("rank ", rank, data)
-sys.stdout.flush()
+print("rank ", rank, py_data)
+stdout.flush()
 
 comm.Scatter(data, recv_buf, root=0)
+new_data = comm.scatter(py_data, root=0)
 
 comm.Barrier()
 if rank == 0:
     print()
     print("Final data")
-    sys.stdout.flush()
+
+stdout.flush()
 comm.Barrier()
 
 print("rank ", rank, recv_buf)
+print("rank ", rank, new_data)