update cat-grid-jobs to work with zdab-cat

1 files changed, 11 insertions, 6 deletions

diff --git a/utils/cat-grid-jobs b/utils/cat-grid-jobs
index f71205d..13b3d46 100755
--- a/utils/cat-grid-jobs
+++ b/utils/cat-grid-jobs
@@ -38,7 +38,7 @@ if __name__ == '__main__':
     import argparse
     import matplotlib.pyplot as plt
     import numpy as np
-    from subprocess import check_output
+    import subprocess
     from os.path import join
     import os
     import sys
@@ -78,24 +78,29 @@ if __name__ == '__main__':
     # Note: We send stderr to /dev/null since there can be a lot of warnings
     # about PMT types and fit results
     with open(os.devnull, 'w') as f:
-        output = check_output([join(args.dir,"zdab-cat"),args.zdab],stderr=f)
-
-    data = yaml.load(output,Loader=Loader)
+        popen = subprocess.Popen([join(args.dir,"zdab-cat"),args.zdab],stdout=subprocess.PIPE,stderr=f)
 
     total_events = 0
     events_with_fit = 0
 
-    for i, event in enumerate(data['data']):
-        for ev in event['ev']:
+    for data in yaml.load_all(popen.stdout,Loader=Loader):
+        if 'ev' not in data:
+            continue
+
+        for ev in data['ev']:
             run = ev['run']
             gtid = ev['gtid']
 
+            print((run,gtid))
+
             if (run,gtid) in fit_results:
                 ev['fit'] = fit_results[(run,gtid)]
                 events_with_fit += 1
 
             total_events += 1
 
+    popen.wait()
+
     # Print out number of fit results that were added. Hopefully, this will
     # make it easy to catch an error if, for example, this gets run with a
     # mismatching zdab and fit results