loop over runs in dc, dc-closure-test, and plot-dc to prevent using too much memory

1 files changed, 12 insertions, 1 deletions

diff --git a/utils/plot-dc b/utils/plot-dc
index 8048192..1ecbbf9 100755
--- a/utils/plot-dc
+++ b/utils/plot-dc
@@ -103,13 +103,24 @@ if __name__ == '__main__':
     parser = argparse.ArgumentParser("plot fit results")
     parser.add_argument("filenames", nargs='+', help="input files")
     parser.add_argument("--save", action='store_true', default=False, help="save corner plots for backgrounds")
+    parser.add_argument("--nhit-thresh", type=int, default=None, help="nhit threshold to apply to events before processing (should only be used for testing to speed things up)")
     args = parser.parse_args()
 
     setup_matplotlib(args.save)
 
     import matplotlib.pyplot as plt
 
-    ev = get_events(args.filenames, merge_fits=True)
+    # Loop over runs to prevent using too much memory
+    evs = []
+    data_filenames = [filename for filename in args.filenames if 'SNO' in filename]
+    mc_filenames = [filename for filename in args.filenames if 'SNO' not in filename]
+    if len(data_filenames):
+        rhdr = pd.concat([read_hdf(filename, "rhdr").assign(filename=filename) for filename in data_filenames],ignore_index=True)
+        for run, df in rhdr.groupby('run'):
+            evs.append(get_events(df.filename.values, merge_fits=True, nhit_thresh=args.nhit_thresh))
+    if len(mc_filenames):
+        evs.append(get_events(mc_filenames, merge_fits=True, nhit_thresh=args.nhit_thresh))
+    ev = pd.concat(evs)
 
     ev = ev[ev.prompt & ~np.isnan(ev.fmin)]