diff options
-rwxr-xr-x | utils/chi2 | 19 |
1 files changed, 17 insertions, 2 deletions
@@ -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 |