Reduce a list with operator as summing each element in mpi4py

Question

I am writing a mpi python code. For example, four procs have data below:

data on procs0: [1, 2, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0]
data on procs1: [0, 0, 0, 4, 5, 6, 0, 0, 0, 0, 0, 0]
data on procs2: [0, 0, 0, 0, 0, 0, 7, 8, 9, 0, 0, 0]
data on procs3: [0, 0, 0, 0, 0, 0, 0, 0, 0, 10, 11, 12]

I want to use reduce function in mpi4py library to reduce the data on procs0, with the result:

result on procs0: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]

How can I make it using mpi4py lib function?

edited: above is a simple special case, set can not be used, see another case below:

data on procs0: [1,0,0,0,0,0]
data on procs1: [0,2,0,0,0,0]
data on procs2: [0,0,0,3,0,0]
data on procs3: [0,0,0,0,4,5]

the ideal result must be:

result on procs0: [1,2,0,3,4,5]

Answer 1

I'm not sure from your question if you want a sum of the data, or the max. I've written up a simple example using the mpi Reduce function, which computes the sum.

#!/usr/bin/env python
import numpy as np
from mpi4py import MPI
comm = MPI.COMM_WORLD

comm.Barrier()
t_start = MPI.Wtime()

# this array lives on each processor
data = np.zeros(5)
for i in xrange(comm.rank, len(data), comm.size):
    # set data in each array that is different for each processor
    data[i] = i

# print out the data arrays for each processor
print '[%i]'%comm.rank, data
comm.Barrier()

# the 'totals' array will hold the sum of each 'data' array
if comm.rank==0:
    # only processor 0 will actually get the data
    totals = np.zeros_like(data)
else:
    totals = None

# use MPI to get the totals 
comm.Reduce(
    [data, MPI.DOUBLE],
    [totals, MPI.DOUBLE],
    op = MPI.SUM,
    root = 0
)

# print out the 'totals'
# only processor 0 actually has the data
print '[%i]'%comm.rank, totals

comm.Barrier()
t_diff = MPI.Wtime() - t_start
if comm.rank==0: print t_diff

Saving this code in the file reduce_test.py and running it with the command mpirun -np 3 ./reduce_test.py gave the following output on my machine:

[0] [ 0.  0.  0.  3.  0.]
[1] [ 0.  1.  0.  0.  4.]
[2] [ 0.  0.  2.  0.  0.]
[1] None
[2] None
[0] [ 0.  1.  2.  3.  4.]
0.00260496139526

Note that changing the parameter op = MPI.SUM in the call to comm.Reduce to op = MPI.MAX will compute the maximum instead of the sum.