speed up nll function in chi2

1 files changed, 17 insertions, 2 deletions

diff --git a/utils/chi2 b/utils/chi2
index 277349d..7f24c1f 100755
--- a/utils/chi2
+++ b/utils/chi2
@@ -98,6 +98,17 @@ def get_mc_hists(data,x,bins,apply_norm=False):
 
     Returns a dictionary mapping particle id combo -> histogram.
     """
+    ke_dict = {}
+    for id in (20,22,2020,2022,2222):
+        ke_dict[id] = data[data.id == id].ke.values
+    return get_mc_hists_fast(ke_dict,x,bins,apply_norm)
+
+def get_mc_hists_fast(ke_dict,x,bins,apply_norm=False):
+    """
+    Same as get_mc_hists() but the first argument is a dictionary mapping
+    particle id -> kinetic energy array. This is much faster than selecting the
+    events from the dataframe every time.
+    """
     mc_hists = {}
 
     # FIXME: May need to increase number of bins here
@@ -105,7 +116,7 @@ def get_mc_hists(data,x,bins,apply_norm=False):
     bincenters2 = (bins2[1:] + bins2[:-1])/2
 
     for id in (20,22,2020,2022,2222):
-        ke = data[data.id == id].ke.values
+        ke = ke_dict[id]
 
         if id in (20,2020):
             ke = ke*x[1]
@@ -159,6 +170,10 @@ def make_nll(data, muons, mc, bins):
     data_hists = get_data_hists(data,bins)
     muon_hists = get_data_hists(muons,bins)
 
+    ke_dict = {}
+    for id in (20,22,2020,2022,2222):
+        ke_dict[id] = data[data.id == id].ke.values
+
     def nll(x, grad=None):
         if any(x[i] < 0 for i in range(len(x))):
             return np.inf
@@ -195,7 +210,7 @@ def make_nll(data, muons, mc, bins):
         #
         # Also, it's critical that I first adjust the data energy by whatever
         # amount I find with the stopping muons and Michel distributions.
-        mc_hists = get_mc_hists(mc,x,bins,apply_norm=True)
+        mc_hists = get_mc_hists_fast(ke_dict,x,bins,apply_norm=True)
 
         # Calculate the negative log of the likelihood of observing the data
         # given the fit parameters