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

3 files changed, 28 insertions, 3 deletions

diff --git a/utils/dc b/utils/dc
index 6c7d86a..40ae13e 100755
--- a/utils/dc
+++ b/utils/dc
@@ -273,13 +273,20 @@ if __name__ == '__main__':
     parser.add_argument("--steps", type=int, default=100000, help="number of steps in the MCMC chain")
     parser.add_argument("--save", action="store_true", default=False, help="save plots")
     parser.add_argument("--mc", nargs='+', required=True, help="atmospheric MC files")
+    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 = []
+    rhdr = pd.concat([read_hdf(filename, "rhdr").assign(filename=filename) for filename in args.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))
+    ev = pd.concat(evs)
+
     ev = ev[ev.prompt]
     ev = ev[ev.nhit_cal > 100]
 
diff --git a/utils/dc-closure-test b/utils/dc-closure-test
index ba5daf8..e23f554 100755
--- a/utils/dc-closure-test
+++ b/utils/dc-closure-test
@@ -425,13 +425,20 @@ if __name__ == '__main__':
     parser.add_argument("--save", action="store_true", default=False, help="save plots")
     parser.add_argument("--mc", nargs='+', required=True, help="atmospheric MC files")
     parser.add_argument("-n", type=int, default=10, help="number of fits to run")
+    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 = []
+    rhdr = pd.concat([read_hdf(filename, "rhdr").assign(filename=filename) for filename in args.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))
+    ev = pd.concat(evs)
+
     ev = ev[ev.prompt]
     ev = ev[ev.nhit_cal > 100]
 
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)]