a fitter class to handle mass fit

.gitignore

@@ -0,0 +1,4 @@
+*.root
+*.pyc
+*.pdf
+*.png

fit/roofit_helper.py

@@ -0,0 +1,120 @@
+import ROOT as rt
+from math import sqrt
+
+
+def CMS_lumi():
+    mark = rt.TLatex()
+    mark.SetNDC()
+    lumistamp = '2018 (13 TeV)'
+    fontScale = 1.0
+    cmsTextSize = 0.042 * fontScale * 1.25
+    extraOverCmsTextSize  = 0.76
+    extraTextSize = extraOverCmsTextSize*cmsTextSize
+
+    mark.SetTextAlign(11)
+    mark.SetTextSize(cmsTextSize)
+    mark.SetTextFont(61)
+    mark.DrawLatex(rt.gPad.GetLeftMargin(), 1 - (rt.gPad.GetTopMargin() - 0.017), "CMS")
+    mark.SetTextSize(0.042 * fontScale)
+    mark.SetTextFont(52)
+    mark.DrawLatex(rt.gPad.GetLeftMargin() + 0.09, 1 - (rt.gPad.GetTopMargin() - 0.017),  "Preliminary")
+    mark.SetTextSize(extraTextSize)
+    mark.SetTextFont(42)
+    mark.SetTextAlign(31)
+    mark.DrawLatex(1 - rt.gPad.GetRightMargin(), 1 - ( rt.gPad.GetTopMargin() - 0.017), lumistamp)
+    return mark
+
+
+def canvas_create(xframe,xmin,xmax,nbin,xtitle,boundYto0=True):
+    c2 = rt.TCanvas('fig_binnedFit', 'fit', 800, 600)
+    c2.SetGrid()
+    rt.gPad.SetLeftMargin(0.12)
+    rt.gPad.SetRightMargin(0.05)
+    rt.gPad.SetBottomMargin(0.15)
+    #xframe.GetYaxis().SetTitle("Events / {0:.0f} MeV".format((xmax - xmin)/nbin*1000.))
+    xframe.GetXaxis().SetTitle(xtitle)
+    xframe.SetStats(0)
+    if (boundYto0): xframe.SetMinimum(0)
+    xframe.GetYaxis().SetLabelSize(0.045)
+    xframe.GetYaxis().SetLabelOffset(0.007)
+    xframe.GetYaxis().SetTitleSize(0.06)
+    xframe.GetYaxis().SetTitleOffset(0.90)
+    xframe.GetXaxis().SetLabelSize(0.045)
+    xframe.GetXaxis().SetLabelOffset(0.007)
+    xframe.GetXaxis().SetTitleSize(0.06)
+    xframe.GetXaxis().SetTitleOffset(0.95)
+    xframe.SetTitle("")
+    xframe.Draw()
+    return c2
+
+def canvas_create_pull(xframe,xmin,xmax,nbin,xtitle,var,boundYto0=True):
+    c2 = rt.TCanvas('fig_binnedFit', 'fit', 800, 800)
+    c2.cd()
+    c2.SetGrid()
+    top = rt.TPad("top", "top", 0., 0.25, 1., 1.)
+    top.SetGrid()
+    top.SetLeftMargin(0.12)
+    top.SetRightMargin(0.05)
+    top.SetBottomMargin(0.0)
+    top.Draw()
+    top.cd()
+    #xframe.GetYaxis().SetTitle("Events / {0:.0f} MeV".format((xmax - xmin)/nbin*1000.))
+    #xframe.GetXaxis().SetTitle(xtitle)
+    xframe.SetStats(0)
+    if (boundYto0): xframe.SetMinimum(0)
+    xframe.GetYaxis().SetLabelSize(0.045)
+    xframe.GetYaxis().SetLabelOffset(0.007)
+    xframe.GetYaxis().SetTitleSize(0.06)
+    xframe.GetYaxis().SetTitleOffset(0.90)
+    xframe.GetXaxis().SetLabelSize(0.045)
+    xframe.GetXaxis().SetLabelOffset(0.007)
+    xframe.GetXaxis().SetTitleSize(0.06)
+    xframe.GetXaxis().SetTitleOffset(0.95)
+    xframe.SetTitle("")
+    xframe.Draw()
+
+    hpull = xframe.pullHist()
+    xframe3 = var.frame()
+    xframe3.addPlotable(hpull,"P")
+
+    scale = 3.0
+    c2.cd()
+    bottom = rt.TPad("bottom", "bottom", 0., 0., 1., 0.25)
+    bottom.SetGrid()
+    bottom.SetLeftMargin(0.12)
+    bottom.SetRightMargin(0.05)
+    bottom.SetBottomMargin(0.4)
+    bottom.SetTopMargin(0.03)
+    bottom.Draw()
+    bottom.cd()
+    xframe3.GetXaxis().SetTitle(xtitle)
+    xframe3.SetStats(0)
+    xframe3.GetYaxis().SetLabelSize(0.045*scale)
+    xframe3.GetYaxis().SetLabelOffset(0.007)
+    xframe3.GetYaxis().SetTitleSize(0.06*scale)
+    xframe3.GetYaxis().SetTitleOffset(0.90/scale)
+    xframe3.GetXaxis().SetLabelSize(0.045*scale)
+    xframe3.GetXaxis().SetLabelOffset(0.007)
+    xframe3.GetXaxis().SetTitleSize(0.06*scale)
+    xframe3.GetXaxis().SetTitleOffset(0.95)
+    xframe3.SetTitle("")
+    xframe3.GetYaxis().SetTitle("Pull")
+    xframe3.SetMinimum(-3.5)
+    xframe3.SetMaximum(3.5)
+
+    xframe3.Draw()
+
+    return c2, top, bottom
+
+def pt_create(mva,nsig,nsigError,nbkg):
+  pt = rt.TPaveText(0.72,0.37,0.92,0.63,"brNDC")
+  pt.SetFillColor(0)
+  pt.SetBorderSize(1)
+  pt.SetTextFont(42);
+  pt.SetTextSize(0.04);
+  pt.SetTextAlign(12)
+  pt.AddText("MVA cut: {0}".format(mva))
+  pt.AddText("S: {0:.1f}#pm{1:.1f}".format(nsig,nsigError))
+  pt.AddText("B: {0:.1f}".format(nbkg))
+  pt.AddText("S/#sqrt{{S+B}}: {0:.2f}".format(nsig/sqrt(nsig + nbkg)))
+  return pt

fit/roofit_models.py

@@ -0,0 +1,23 @@
+root_function_DoubleSidedCB="""
+
+double DoubleSidedCB2(double x, double mu, double width, double a1, double p1, double a2, double p2)
+{
+  double u   = (x-mu)/width;
+  double A1  = TMath::Power(p1/TMath::Abs(a1),p1)*TMath::Exp(-a1*a1/2);
+  double A2  = TMath::Power(p2/TMath::Abs(a2),p2)*TMath::Exp(-a2*a2/2);
+  double B1  = p1/TMath::Abs(a1) - TMath::Abs(a1);
+  double B2  = p2/TMath::Abs(a2) - TMath::Abs(a2);
+
+  double result(1);
+  if      (u<-a1) result *= A1*TMath::Power(B1-u,-p1);
+  else if (u<a2)  result *= TMath::Exp(-u*u/2);
+  else            result *= A2*TMath::Power(B2+u,-p2);
+  return result;
+}
+
+double ROOT_DoubleSidedCB(double* x, double *par)
+{
+  return(par[0] * DoubleSidedCB2(x[0], par[1],par[2],par[3],par[4],par[5],par[6]));
+}
+"""
+

fit/run_fit_KEE_lowq2_bparkPU.py

@@ -0,0 +1,173 @@
+import uproot
+import pandas as pd
+import numpy as np
+from collections import OrderedDict, defaultdict
+from scipy import interp
+from rootpy.io import root_open
+from rootpy.plotting import Hist
+from root_numpy import fill_hist, array2root, array2tree
+from root_pandas import to_root
+import itertools
+import PyPDF2
+#import makePlot_fitPeak_unbinned as fit_unbinned
+import os, sys, copy
+import xgboost as xgb
+
+LOWQ2_LOW = 1.05
+LOWQ2_UP = 2.45
+JPSI_LOW = 2.8
+JPSI_UP = 3.25
+PSI2S_LOW = 3.45
+PSI2S_UP = 3.8
+HIGHQ2_LOW = 4.0
+HIGHQ2_UP = 4.87
+
+def get_df(root_file_name, tree='mytreefit', branches=['*']):
+    print('Opening file {}...'.format(root_file_name))
+    f = uproot.open(root_file_name)
+    if len(f.allkeys()) == 0:
+        return pd.DataFrame()
+    print('Not an null file')
+    #df = uproot.open(root_file_name)["tree"].pandas.df()
+    #df = pd.DataFrame(uproot.open(root_file_name)["tree"].arrays(namedecode="utf-8"))
+    df = pd.DataFrame(f[tree].arrays(branches=branches))
+    print('Finished opening file {}...'.format(root_file_name))
+    return df
+
+if __name__ == "__main__":
+  eleType = 'pf'
+  log = 'log_kee_bparkPU_v7.3_{}_nonreg.csv'.format(eleType)
+  info = defaultdict(dict)
+
+  nparts = range(8)
+
+  '''
+  info['pf']['inputfile'] = '../data/data_PFe_v7.3/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_PFe_v7.3_data_mvaCut0.root'
+  info['pf']['nonresonant_mc'] = '../data/data_PFe_v7.3/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_PFe_v7.3_{}_MC.root'.format('marker')
+  info['pf']['jpsi_mc'] = '../data/data_PFe_v7.3/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_PFe_v7.3_{}_MCres.root'.format('marker')
+  info['pf']['isgn'] = '../data/data_PFe_v7.3/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_PFe_v7.3_0_MC.root'
+  info['pf']['ikjpsi'] = '../data/data_PFe_v7.3/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_PFe_v7.3_0_MCres.root'
+  info['pf']['ikstaree'] = '../data/data_PFe_v7.3/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_PFe_v7.3_0_MC_kstaree.root'
+  #info['pf']['ikstaree'] = '../data/data_PFe_v7.3/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_PFe_v7.3_0_MC_kstarjpsi.root'
+  info['pf']['jpsi_mva_wp'] = 6.0
+  info['pf']['n_data_jpsi'] = 7069.5
+  info['pf']['n_mc_jpsi'] = 563421.0
+  info['pf']['n_mc_lowq2'] = 617615.0
+  '''
+  '''
+  info['mix']['inputfile'] = '../data/data_LowPtPF_v7.3/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_LowPtPF_v7.3_data_mvaCut0.root'
+  info['mix']['nonresonant_mc'] = '../data/data_LowPtPF_v7.3/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_LowPtPF_v7.3_{}_MC.root'.format('marker')
+  info['mix']['jpsi_mc'] = '../data/data_LowPtPF_v7.3/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_LowPtPF_v7.3_{}_MCres.root'.format('marker')
+  info['mix']['isgn'] = '../data/data_LowPtPF_v7.3/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_LowPtPF_v7.3_0_MC.root'
+  info['mix']['ikjpsi'] = '../data/data_LowPtPF_v7.3/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_LowPtPF_v7.3_0_MCres.root'
+  info['mix']['ikstaree'] = '../data/data_LowPtPF_v7.3/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_LowPtPF_v7.3_0_MC_kstaree.root'
+  info['mix']['jpsi_mva_wp'] = 5.0
+  info['mix']['n_data_jpsi'] = 5533.45
+  info['mix']['n_mc_jpsi'] = 563421.0
+  info['mix']['n_mc_lowq2'] = 617615.0
+  '''
+  '''
+  info['pf']['inputfile'] = '../data/data_PFe_v7.3_nonreg_rmVar/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_PFe_v7.3_nonreg_rmVar_data_mvaCut0.root'
+  info['pf']['nonresonant_mc'] = '../data/data_PFe_v7.3_nonreg_rmVar/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_PFe_v7.3_nonreg_rmVar_{}_MC.root'.format('marker')
+  info['pf']['jpsi_mc'] = '../data/data_PFe_v7.3_nonreg_rmVar/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_PFe_v7.3_nonreg_rmVar_{}_MCres.root'.format('marker')
+  info['pf']['isgn'] = '../data/data_PFe_v7.3_nonreg_rmVar/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_PFe_v7.3_nonreg_rmVar_0_MC.root'
+  info['pf']['ikjpsi'] = '../data/data_PFe_v7.3_nonreg_rmVar/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_PFe_v7.3_nonreg_rmVar_0_MCres.root'
+  info['pf']['ikstaree'] = '../data/data_PFe_v7.3/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_PFe_v7.3_0_MC_kstaree.root'
+  #info['pf']['jpsi_mva_wp'] = 5.0
+  #info['pf']['n_data_jpsi'] = 6853.0
+  info['pf']['jpsi_mva_wp'] = 6.5
+  info['pf']['n_data_jpsi'] = 5611.69
+
+  #info['pf']['n_mc_jpsi'] = 563421.0
+  #info['pf']['n_mc_lowq2'] = 617615.0
+  info['pf']['n_mc_jpsi'] = 209954.0
+  info['pf']['n_mc_lowq2'] = 231498.0
+  '''
+
+  info['pf']['inputfile'] = '../data/data_PFe_v7.3_nonreg_ottoCut/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_PFe_v7.3_nonreg_ottoCut_data_mvaCut0.root'
+  info['pf']['nonresonant_mc'] = '../data/data_PFe_v7.3_nonreg_ottoCut/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_PFe_v7.3_nonreg_ottoCut_{}_MC.root'.format('marker')
+  info['pf']['jpsi_mc'] = '../data/data_PFe_v7.3_nonreg_ottoCut/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_PFe_v7.3_nonreg_ottoCut_{}_MCres.root'.format('marker')
+  info['pf']['isgn'] = '../data/data_PFe_v7.3_nonreg_ottoCut/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_PFe_v7.3_nonreg_ottoCut_0_MC.root'
+  info['pf']['ikjpsi'] = '../data/data_PFe_v7.3_nonreg_ottoCut/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_PFe_v7.3_nonreg_ottoCut_0_MCres.root'
+  info['pf']['ikstaree'] = '../data/data_PFe_v7.3/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_PFe_v7.3_0_MC_kstaree.root'
+  info['pf']['jpsi_mva_wp'] = 8.3
+  info['pf']['n_data_jpsi'] = 4779.12
+  info['pf']['n_mc_jpsi'] = 211912.0
+  info['pf']['n_mc_lowq2'] = 232115.0
+
+  info['mix']['inputfile'] = '../data/data_LowPtPF_v7.3_nonreg/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_LowPtPF_v7.3_nonreg_data_mvaCut0.root'
+  info['mix']['nonresonant_mc'] = '../data/data_LowPtPF_v7.3_nonreg/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_LowPtPF_v7.3_nonreg_{}_MC.root'.format('marker')
+  info['mix']['jpsi_mc'] = '../data/data_LowPtPF_v7.3_nonreg/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_LowPtPF_v7.3_nonreg_{}_MCres.root'.format('marker')
+  info['mix']['isgn'] = '../data/data_LowPtPF_v7.3_nonreg/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_LowPtPF_v7.3_nonreg_0_MC.root'
+  info['mix']['ikjpsi'] = '../data/data_LowPtPF_v7.3_nonreg/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_LowPtPF_v7.3_nonreg_0_MCres.root'
+  info['mix']['ikstaree'] = '../data/data_LowPtPF_v7.3_nonreg/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_LowPtPF_v7.3_nonreg_0_MC_kstaree.root'
+  info['mix']['jpsi_mva_wp'] = 8.6
+  info['mix']['n_data_jpsi'] = 2084.77
+  info['mix']['n_mc_jpsi'] = 211912.0
+  info['mix']['n_mc_lowq2'] = 232115.0
+  info['mix']['ikstaree'] = '../data/data_PFe_v7.3/forMeas_xgbmodel_kee_12B_kee_correct_pu_Depth17_PFe_v7.3_0_MC_kstaree.root'
+
+  num_mctoys = None #5000
+
+  rk_electron = 7.20e-3
+
+  jpsi_mva_wp = info[eleType]['jpsi_mva_wp']
+  n_data_jpsi = info[eleType]['n_data_jpsi']
+
+  selection = {}
+
+  selection['jpsi'] = '(Mll > 2.9) and (Mll < 3.2)'
+  selection['psi2s'] = '(Mll > 3.55) and (Mll < 3.8)'
+  selection['lowq2'] = '(Mll > 1.05) and (Mll < 2.45)'
+  selection['highq2'] = '(Mll > @PSI2S_LOW) and (Mll < @PSI2S_UP)'
+
+  selection['Dmass'] = '(KLmassD0 > 2.0)'
+  #selection['elePt'] = '(L1pt > 10) and (L2pt > 10)'
+  #selection['hlt'] = '(Mu9_IP6 == True)'
+
+  mc_branches = ['Bmass', 'Mll', 'KLmassD0', 'xgb'] #+['Mu9_IP6']#+ ['L1pt', 'L2pt'] 
+
+  jpsi_mc_branches = [get_df(info[eleType]['jpsi_mc'].replace('marker', str(i)), branches=mc_branches) for i in nparts]
+  nonresonant_mc_branches = [get_df(info[eleType]['nonresonant_mc'].replace('marker', str(i)), branches=mc_branches) for i in nparts]
+
+  eff = {}
+  # efficieny of preselection and di-electron type (after nanoAOD preselection)
+  eff['jpsi_presel'] = float(jpsi_mc_branches[0].shape[0]) / info[eleType]['n_mc_jpsi']
+  eff['lowq2_presel'] = float(nonresonant_mc_branches[0].shape[0]) / info[eleType]['n_mc_lowq2']
+  # efficiency of q2
+  eff['jpsi_q2'] = float(jpsi_mc_branches[0].query(selection['jpsi']).shape[0]) / float(jpsi_mc_branches[0].shape[0])
+  eff['lowq2_q2'] = float(nonresonant_mc_branches[0].query(selection['lowq2']).shape[0]) / float(nonresonant_mc_branches[0].shape[0])
+  # efficiency of bdt of jpsi
+  eff['jpsi_bdt'] = np.mean([float(jpsi_mc_branches[i].query(' and '.join([selection['jpsi'], selection['Dmass'], '(xgb > {})'.format(jpsi_mva_wp)])).shape[0]) / float(jpsi_mc_branches[i].query(selection['jpsi']).shape[0]) for i in nparts])
+  #eff['jpsi_bdt'] = np.mean([float(jpsi_mc_branches[i].query(' and '.join([selection['jpsi'], selection['Dmass'], selection['hlt'], '(xgb > {})'.format(jpsi_mva_wp)])).shape[0]) / float(jpsi_mc_branches[i].query(selection['jpsi']).shape[0]) for i in nparts])
+
+  if eleType == 'pf':
+    #mvaCut = np.linspace(8.0, 8.7, 8)
+    mvaCut = np.array([8.3, ])
+  else:
+    #mvaCut = np.linspace(8.3, 9.0, 8)
+    mvaCut = np.array([8.6, ])
+
+  for cut in mvaCut:
+    eff_lowq2_bdt = np.mean([float(nonresonant_mc_branches[i].query(' and '.join([selection['lowq2'], '(xgb > @cut)', selection['Dmass']])).shape[0]) / float(nonresonant_mc_branches[i].query(selection['lowq2']).shape[0]) for i in nparts])
+    #eff_lowq2_bdt = np.mean([float(nonresonant_mc_branches[i].query(' and '.join([selection['lowq2'], selection['hlt'], '(xgb > @cut)', selection['Dmass']])).shape[0]) / float(nonresonant_mc_branches[i].query(selection['lowq2']).shape[0]) for i in nparts])
+
+    expected_signal = rk_electron * n_data_jpsi * eff['lowq2_presel'] * eff['lowq2_q2'] * eff_lowq2_bdt / (eff['jpsi_presel'] * eff['jpsi_q2'] * eff['jpsi_bdt'])
+
+    if num_mctoys is not None:
+      com = 'python KEE_lowq2_roofit_plb_pull_modified.py -i {0} -o lowq2_{1} --isgn={2} --ikjpsi={3} --sel_primitive="sgn,bkg_kjpsi" --fit_primitive --mvacut={4} --set_expected_sgn={5} --number_of_mctoys={6} --log={7}'.format(info[eleType]['inputfile'], eleType, info[eleType]['isgn'], info[eleType]['ikjpsi'], cut, expected_signal, num_mctoys, log)
+    else:
+      com = 'python KEE_lowq2_roofit_plb_pull_modified_kde.py -i {0} -o lowq2_{1} --isgn={2} --ikjpsi={3} --ikstaree={4} --sel_primitive="sgn,bkg_kjpsi,bkg_kstaree" --fit_primitive --mvacut={5} --set_expected_sgn={6} --log={8}'.format(info[eleType]['inputfile'], eleType, info[eleType]['isgn'], info[eleType]['ikjpsi'], info[eleType]['ikstaree'], cut, expected_signal, num_mctoys, log)
+      #com = 'python KEE_lowq2_roofit_plb_pull_modified_kde.py -i {0} -o lowq2_{1} --isgn={2} --ikjpsi={3} --sel_primitive="sgn,bkg_kjpsi" --fit_primitive --mvacut={4} --set_expected_sgn={5} --log={7}'.format(info[eleType]['inputfile'], eleType, info[eleType]['isgn'], info[eleType]['ikjpsi'], cut, expected_signal, num_mctoys, log)
+
+    os.system(com)
+    print('cut: {} \n \t expected signal: {} \n \t lowq2 bdt eff: {}'.format(cut, expected_signal, eff_lowq2_bdt))
+    print('\t {}, {}'.format(eff['lowq2_presel'], eff['lowq2_q2']))
+    print('\t overall eff: {}'.format(eff['lowq2_presel'] * eff['lowq2_q2'] * eff_lowq2_bdt))
+
+  for key, value in eff.items():
+    print('{}: {}'.format(key, value))
+
+
+
+