This example can be found in /sdf/group/lcls/ds/ana/tutorials/psana1_examples/mpiGatherv.py. It uses the more performant, but more difficult to use, "Gatherv" array implementation (vs. the lowercase "gather" which uses python objects shown here). There is also "Gather" method which more easily gathers arrays from all ranks that are of the same size.
from psana import * import numpy as np from mpi4py import MPI comm = MPI.COMM_WORLD rank = comm.Get_rank() size = comm.Get_size() ds = DataSource("exp=xpptut15:run=54:smd") det = Detector('cspad',ds.env()) intensities = [] # use a list instead of array because faster to append for nevent,evt in enumerate(ds.events()): if nevent==5: break if nevent%size!=rank: continue intensities.append(det.raw(evt).sum()) # real Gatherv logic starts here lengths = np.array(comm.gather(len(intensities))) # get list of lengths mysend = np.array(intensities) print(f'Rank: {rank}, sending: {mysend}') myrecv = None if rank==0: myrecv = np.empty((sum(lengths)),mysend.dtype) # allocate receive buffer comm.Gatherv(sendbuf=mysend, recvbuf=[myrecv, lengths]) if rank==0: start = 0 # look in the receive buffer for the contribution from each rank for r,mylen in enumerate(lengths): print (f'Rank 0 received {myrecv[start:start+mylen]} from rank {r}') start += mylen
Overview
Content Tools