-
Notifications
You must be signed in to change notification settings - Fork 1
/
Copy pathbatchtest_mpi.py
61 lines (51 loc) · 1.59 KB
/
batchtest_mpi.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
import psana
import numpy as np
import sys
import socket
import os
from mpi4py import MPI
comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('-d','--dir',nargs='+', help="list of directories",required=True)
args = parser.parse_args()
for d in args.dir:
nerror=0
if not os.path.isfile(d+'/e307-r0999-s00-c00.xtc'):
print('*** Error: Host',socket.gethostname(),'cannot access directory',d)
nerror+=1
if nerror:
sys.exit()
for d in args.dir:
firsttime = True
ds = psana.DataSource('dir=%s:exp=xcstut13:run=999:idx'%d)
src = psana.Source('DetInfo(XcsBeamline.0:Princeton.0)')
maxEventsPerNode=2
for run in ds.runs():
times = run.times()
mylength = len(times)/size
if mylength>maxEventsPerNode: mylength=maxEventsPerNode
mytimes= times[rank*mylength:(rank+1)*mylength]
for i in range(mylength):
evt = run.event(mytimes[i])
if evt is None:
print('*** event fetch failed')
continue
cam = evt.get(psana.Princeton.FrameV1,src)
if cam is None:
print('*** failed to get cam')
continue
if firsttime:
sum=cam.data().astype(np.float)
firsttime=False
else:
sum+=cam.data()
sumall = np.empty_like(sum)
#sum the images across mpi cores
comm.Reduce(sum,sumall)
if rank==0:
print('dir:',d)
print('sum:',np.sum(sumall))
MPI.Finalize()