speed up chi2 and dm-search

2 files changed, 19 insertions, 11 deletions

diff --git a/utils/chi2 b/utils/chi2
index a98986a..117d0a0 100755
--- a/utils/chi2
+++ b/utils/chi2
@@ -206,13 +206,15 @@ def get_data_hists(data,bins,scale=1.0):
     return data_hists
 
 def make_nll(data, muons, mc, atmo_scale_factor, muon_scale_factor, bins, reweight=False, print_nll=False):
-    df_dict = {}
+    df_dict = dict(tuple(mc.groupby('id')))
     for id in (20,22,2020,2022,2222):
-        df_dict[id] = mc[mc.id == id]
+        if id not in df_dict:
+            df_dict[id] = mc.iloc[:0]
 
-    df_dict_muon = {}
+    df_dict_muon = dict(tuple(muons.groupby('id')))
     for id in (20,22,2020,2022,2222):
-        df_dict_muon[id] = muons[muons.id == id]
+        if id not in df_dict_muon:
+            df_dict_muon[id] = muons.iloc[:0]
 
     data_hists = get_data_hists(data,bins)
 
@@ -400,9 +402,11 @@ def do_fit(data,muon,data_mc,weights,atmo_scale_factor,muon_scale_factor,bins,st
     # Get the total number of "universes" simulated in the GENIE reweight tool
     nuniverses = weights['universe'].max()+1
 
+    weights_dict = dict(tuple(weights.groupby('universe')))
+
     nlls = []
     for universe in range(nuniverses):
-        data_mc_with_weights = pd.merge(data_mc,weights[weights.universe == universe],how='left',on=['run','evn'])
+        data_mc_with_weights = pd.merge(data_mc,weights_dict[universe],how='left',on=['run','evn'])
         data_mc_with_weights.weight = data_mc_with_weights.weight.fillna(1.0)
 
         nll = make_nll(data,muon,data_mc_with_weights,atmo_scale_factor,muon_scale_factor,bins,reweight=True,print_nll=print_nll)
diff --git a/utils/dm-search b/utils/dm-search
index 34509ac..ed05a91 100755
--- a/utils/dm-search
+++ b/utils/dm-search
@@ -223,13 +223,15 @@ def get_data_hists(data,bins,scale=1.0):
     return data_hists
 
 def make_nll(dm_particle_id, dm_mass, dm_energy, data, muons, mc, atmo_scale_factor, muon_scale_factor, bins, reweight=False, print_nll=False, dm_sample=None):
-    df_dict = {}
+    df_dict = dict(tuple(mc.groupby('id')))
     for id in (20,22,2020,2022,2222):
-        df_dict[id] = mc[mc.id == id]
+        if id not in df_dict:
+            df_dict[id] = mc.iloc[:0]
 
-    df_dict_muon = {}
+    df_dict_muon = dict(tuple(muons.groupby('id')))
     for id in (20,22,2020,2022,2222):
-        df_dict_muon[id] = muons[muons.id == id]
+        if id not in df_dict_muon:
+            df_dict_muon[id] = muons.iloc[:0]
 
     data_hists = get_data_hists(data,bins)
 
@@ -341,9 +343,11 @@ def do_fit(dm_particle_id,dm_mass,dm_energy,data,muon,data_mc,weights,atmo_scale
     # Get the total number of "universes" simulated in the GENIE reweight tool
     nuniverses = weights['universe'].max()+1
 
+    weights_dict = dict(tuple(weights.groupby('universe')))
+
     nlls = []
     for universe in range(nuniverses):
-        data_mc_with_weights = pd.merge(data_mc,weights[weights.universe == universe],how='left',on=['run','evn'])
+        data_mc_with_weights = pd.merge(data_mc,weights_dict[universe],how='left',on=['run','evn'])
         data_mc_with_weights.weight = data_mc_with_weights.weight.fillna(1.0)
 
         nll = make_nll(dm_particle_id,dm_mass,dm_energy,data,muon,data_mc_with_weights,atmo_scale_factor,muon_scale_factor,bins,reweight=True,print_nll=print_nll,dm_sample=dm_sample)
@@ -505,7 +509,7 @@ def get_limits(dm_masses,data,muon,data_mc,atmo_scale_factor,muon_scale_factor,b
             # increase our discovery threshold to account for the look
             # elsewhere effect.
             threshold = DISCOVERY_P_VALUE/(len(bins)-1)
-            discovery = poisson.ppf(1-threshold,n) + 1 - n
+            discovery = poisson.ppf(1-threshold,n) - n
             # Here, we scale the discovery threshold by the fraction of the
             # dark matter hist in the histogram range. The idea is that if only
             # a small fraction of the dark matter histogram falls into the