Adding Arrays With "Reduce"

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

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',ds.env())

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()