go back to checking rhdr instead of ev

2 files changed, 14 insertions, 10 deletions

diff --git a/utils/chi2 b/utils/chi2
index fe602ec..b402a4b 100755
--- a/utils/chi2
+++ b/utils/chi2
@@ -504,12 +504,14 @@ if __name__ == '__main__':
 
     import matplotlib.pyplot as plt
 
+    # Loop over runs to prevent using too much memory
+    rhdr = pd.concat([read_hdf(filename, "rhdr").assign(filename=filename) for filename in args.filenames],ignore_index=True)
+
     if args.run_list is not None:
         run_list = np.genfromtxt(args.run_list)
+        rhdr = rhdr[rhdr.run.isin(run_list)]
 
-    # 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'):
         if args.run_list and run not in run_list:
             continue
@@ -624,13 +626,13 @@ if __name__ == '__main__':
     #muon = ev[ev.muon & ev.prompt & ~ev.atm]
     #muon_atm = ev[ev.muon & ev.prompt & ev.atm]
 
-    if (~ev.run.isin(ev_mc.run)).any():
+    if (~rhdr.run.isin(ev_mc.run)).any():
         print_warning("Error! The following runs have no Monte Carlo: %s" % \
-            np.unique(ev.run[~ev.run.isin(ev_mc.run)].values))
+            np.unique(rhdr.run[~rhdr.run.isin(ev_mc.run)].values))
         sys.exit(1)
 
     if not args.pull and not args.coverage:
-        ev_mc = ev_mc[ev_mc.run.isin(ev.run)]
+        ev_mc = ev_mc[ev_mc.run.isin(rhdr.run)]
 
     data_mc = ev_mc[ev_mc.signal & ev_mc.prompt & ~ev_mc.atm]
     data_atm_mc = ev_mc[ev_mc.signal & ev_mc.prompt & ev_mc.atm]
diff --git a/utils/dm-search b/utils/dm-search
index 63bd0c2..7e53484 100755
--- a/utils/dm-search
+++ b/utils/dm-search
@@ -562,12 +562,14 @@ if __name__ == '__main__':
 
     import matplotlib.pyplot as plt
 
+    rhdr = pd.concat([read_hdf(filename, "rhdr").assign(filename=filename) for filename in args.filenames],ignore_index=True)
+
     if args.run_list is not None:
         run_list = np.genfromtxt(args.run_list)
+        rhdr = rhdr[rhdr.run.isin(run_list)]
 
     # 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'):
         if args.run_list and run not in run_list:
             continue
@@ -690,13 +692,13 @@ if __name__ == '__main__':
     #muon = ev[ev.muon & ev.prompt & ~ev.atm]
     #muon_atm = ev[ev.muon & ev.prompt & ev.atm]
 
-    if (~ev.run.isin(ev_mc.run)).any():
+    if (~rhdr.run.isin(ev_mc.run)).any():
         print_warning("Error! The following runs have no Monte Carlo: %s" % \
-            np.unique(ev.run[~ev.run.isin(ev_mc.run)].values))
+            np.unique(rhdr.run[~rhdr.run.isin(ev_mc.run)].values))
         sys.exit(1)
 
-    if not args.pull and not args.test:
-        ev_mc = ev_mc[ev_mc.run.isin(ev.run)]
+    if not args.pull and not args.coverage:
+        ev_mc = ev_mc[ev_mc.run.isin(rhdr.run)]
 
     data_mc = ev_mc[ev_mc.signal & ev_mc.prompt & ~ev_mc.atm]
     data_atm_mc = ev_mc[ev_mc.signal & ev_mc.prompt & ev_mc.atm]