The commands – scatter and gather

Now we are ready to look at the entire script for our demo problem, the scalar product:

from mpi4py import MPI
import numpy as np

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
nprocessors = comm.Get_size()
import splitarray as spa 

if rank == 0:
    # Here we generate data for the example
    n = 150
    u = 0.1*np.arange(n)
    v = - u
    u_split = spa.split_array(u, nprocessors)
    v_split = spa.split_array(v, nprocessors)
else:
    # On all processor we need variables with these names,
    # otherwise we would get an Exception "Variable not defined" in 
    # the scatter command below
    u_split = None
    v_split = None
# These commands run now on all processors
u_split = comm.scatter(u_split, root=0) # the data is portion wise 
                                        # distributed from root
v_split = comm.scatter(v_split, root=0)
# Each processor computes its part of the scalar product
partial_dot = u_split@v_split
# Each processor reports its result...