from psana import DataSource
import numpy as np
import os

# OPTIONAL callback, usually used for realtime plots
# called back on each SRV node, for every smd.event() call below
def test_callback(data_dict):
    print(data_dict)

# this sets the number of h5 files to write. 1 is sufficient for 120Hz operation
# only needed if you are saving h5.
os.environ['PS_SRV_NODES']='1'
ds = DataSource(exp='tmoc00118', run=123, max_events=100)
# batch_size here specifies how often the dictionary of information
# is sent to the SRV nodes
smd = ds.smalldata(filename='my.h5', batch_size=5, callbacks=[test_callback])
run = next(ds.runs())

# necessary (instead of "None") since some ranks may not receive events
# and the smd.sum() below could fail
arrsum = np.zeros((2), dtype=np.int)
for i,evt in enumerate(run.events()):
    myones = np.ones_like(arrsum)
    smd.event(evt, myfloat=2.0, arrint=myones)
    arrsum += myones

if smd.summary:
    smd.sum(arrsum)
    smd.save_summary({'summary_array' : arrsum}, summary_int=1)
smd.done()