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.

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

for nevent,evt in enumerate(ds.events()):
    if nevent%size==rank:  # different ranks look at different events
        if 'image' not in locals():
            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()