speed up chi2 and dm-search

1 files changed, 9 insertions, 5 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)