Adding Arrays With "Reduce"

This script lives in /reg/g/psdm/tutorials/examplePython/mpiReduce.py and demonstrates how to add arrays distributed across cores (e.g. area detector images) in a parallel MPI job.

Note that this example uses the capitalized "Reduce" method (not "reduce"). This means it works very efficiently with numpy arrays, but is slightly more cumbersome to use.

#run this with: mpirun -n 2 python mpiReduce.py
from psana import *
import numpy as np
ds = DataSource('exp=xpptut15:run=54:smd')
det = Detector('cspad')

from mpi4py import MPI
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

img = None
for nevent,evt in enumerate(ds.events()):
    if nevent%size!=rank: continue # different ranks look at different events
    if img is None:
        img = det.image(evt).astype(np.double)
    else:
        print '*',rank
        img += det.image(evt)
    if nevent>=6: break

imgsum = np.empty_like(img)
comm.Reduce(img,imgsum) # sum the image across all ranks
print 'Rank',rank,'sum:',img.sum()
if rank==0: print 'Total sum:',imgsum.sum()

MPI.Finalize()