import ROOT as r import root_numpy as rn import pickle import numpy as np import matplotlib.pyplot as plt from xgboost import XGBClassifier from tools.data_processing import * from tools.mc_fitter import MC_fit import os import argparse l_flv=['e','mu'] mother_ID=["Ds","Dplus","both"] def select_and_fit(mother_index_fit, l_index, x_cut, iteration, test): #PATHS NEEDED FIT_PATH=l_flv[l_index]+'/fits' BDT_PATH=l_flv[l_index]+'/BDTs/test_'+str(test) with open(BDT_PATH+'/variables_used.pickle', 'rb') as f: branches_needed_0 = pickle.load(f) branches_needed=['cos_thetal']+branches_needed_0 #Data loading MC_Dplus_sig_dict, MC_Ds_sig_dict, _ = load_datasets(l_index) with open('/disk/lhcb_data/davide/Rphipi/NN_for_selection/'+l_flv[l_index]+l_flv[l_index]+'/'+'data_for_NN_'+l_flv[l_index]+l_flv[l_index]+'.pickle', 'rb') as f: data_dict=pickle.load(f) with open('/disk/lhcb_data/davide/Rphipi/MC/for_fit/'+l_flv[l_index]+l_flv[l_index]+'/MC_'+mother_ID[0]+'_'+l_flv[l_index]+l_flv[l_index]+'.pickle', 'rb') as f: MC_Ds_mass_for_fit=pickle.load(f) with open('/disk/lhcb_data/davide/Rphipi/MC/for_fit/'+l_flv[l_index]+l_flv[l_index]+'/MC_'+mother_ID[1]+'_'+l_flv[l_index]+l_flv[l_index]+'.pickle', 'rb') as f: MC_Dplus_mass_for_fit=pickle.load(f) """ Normalising the chi2s in data and MC """ MC_Dplus_sig_dict, MC_Ds_sig_dict, data_dict = norm_chi2(MC_Dplus_sig_dict, MC_Ds_sig_dict, data_dict) """ Applying mass cuts on data, retrieving mass distributions for sig MC and data """ data_mass, mc_mass, data_dict, lower_cut, upper_cut =mass_cut_for_fit(mother_index_fit=mother_index_fit, l_index=l_index, branches_needed=branches_needed, data_dict=data_dict, MC_Dplus_sig_dict=MC_Dplus_sig_dict, MC_Ds_sig_dict=MC_Ds_sig_dict) """ Load the MC DCB shape parameters, if not present fit them """ if os.path.exists(FIT_PATH+'/MC_'+mother_ID[0]+'_CB_params.pickle'): with open(FIT_PATH+'/MC_'+mother_ID[0]+'_CB_params.pickle','rb') as f: Ds_CB_params = pickle.load(f) alpha_Ds_l=Ds_CB_params['alpha_Ds_l'] power_Ds_l=Ds_CB_params['n_Ds_l'] alpha_Ds_r=Ds_CB_params['alpha_Ds_r'] power_Ds_r=Ds_CB_params['n_Ds_r'] frac_Ds=Ds_CB_params['CB fraction'] else: MC_Ds_fit_results, alpha_Ds_l, power_Ds_l, alpha_Ds_r, power_Ds_r, frac_Ds = MC_fit( MC_Ds_mass_for_fit, mother_index=0, l_index=l_index) if os.path.exists(FIT_PATH+'/MC_'+mother_ID[1]+'_CB_params.pickle'): with open(FIT_PATH+'/MC_'+mother_ID[1]+'_CB_params.pickle','rb') as f: Dplus_CB_params = pickle.load(f) alpha_Dplus_l=Dplus_CB_params['alpha_Dplus_l'] power_Dplus_l=Dplus_CB_params['n_Dplus_l'] alpha_Dplus_r=Dplus_CB_params['alpha_Dplus_r'] power_Dplus_r=Dplus_CB_params['n_Dplus_r'] frac_Dplus=Dplus_CB_params['CB fraction'] else: MC_Dplus_fit_results, alpha_Dplus_l, power_Dplus_l, alpha_Dplus_r, power_Dplus_r, frac_Dplus = MC_fit(MC_Dplus_mass_for_fit, mother_index=1, l_index=l_index) """ Preprocess data for XGBoost classifier """ data_2, mc_2, data_mean, mc_mean, data_std, mc_std = preprocess_for_XGBoost(MC_Dplus_sig_dict, MC_Ds_sig_dict, data_dict, branches_needed, mother_index_fit=mother_index_fit) Nsig_from_MC=mc_mass.shape[0] """ BDT selection """ print("BDT cut at {0} fit".format(x_cut)) print("BDT selection..") k = np.random.randint(10) MODEL_PATH=BDT_PATH+'/XG_'+str(k) loaded_model = pickle.load(open(MODEL_PATH+"/XG_"+str(k)+"_.pickle.dat", "rb")) output_XG=loaded_model.predict_proba(data_2) output_XG_signal=loaded_model.predict_proba(mc_2) data_selected=np.array(data_mass[np.where(output_XG[:,1]>x_cut)]) XG_mc_selected=np.array(mc_mass[np.where(output_XG_signal[:,1]>x_cut)]) sig_sel_eff=np.float(XG_mc_selected.shape[0])/np.float(Nsig_from_MC) print("Signal selection efficiency: {0}".format(sig_sel_eff)) """ Select mass window for fit """ cut_indices=[] for i in range(len(data_selected)): if lower_cut<data_selected[i]<upper_cut: cut_indices.append(i) data_cut=data_selected[cut_indices]/1000. lower_cut/=1000. upper_cut/=1000. """ Create a rootfile from which will do the unbinned fit """ data_array=np.array( [data_cut[i] for i in range(len(data_cut))], dtype=[('D_reco_M', np.float32)] ) rn.array2root(data_array, filename='/disk/lhcb_data/davide/Rphipi/selected_data/'+l_flv[l_index]+l_flv[l_index]+'/sel_data_'+l_flv[l_index]+l_flv[l_index]+'.root', treename='decay_tree', mode='recreate', ) f=r.TFile('/disk/lhcb_data/davide/Rphipi/selected_data/'+l_flv[l_index]+l_flv[l_index]+'/sel_data_'+l_flv[l_index]+l_flv[l_index]+'.root') tree=f.Get("decay_tree") #some tricks to pass Roofit a dataset in python mass = np.array([0],dtype=np.float32) branch = tree.GetBranch("D_reco_M") branch.SetAddress(mass) num_entries=tree.GetEntries() m = r.RooRealVar("D mass reco (GeV/c^2)","D mass reco (GeV/c^2)",lower_cut,upper_cut) l = r.RooArgSet(m) data_set = r.RooDataSet("data set", "data set", l) for i in range(num_entries): tree.GetEvent(i) r.RooAbsRealLValue.__assign__(m, mass[0]) data_set.add(l, 1.0) f.Close() """ Retrieve the needed pdfs """ ##Creating D+ signal PDF left tail mean_Dplus= r.RooRealVar("mean_Dplus","mean_Dplus",1.87,1.83,1.91) sigma_Dplus = r.RooRealVar("sigma_Dplus","sigma_Dplus",0.020,0.,0.2) sig_frc_Dplus =r.RooRealVar("Dplus Double crystalball fraction","Dplus Double crystalball fraction",frac_Dplus) al_Dplus_left = r.RooRealVar("alpha_Dplus_left","alpha_Dplus_left",alpha_Dplus_l) pwr_Dplus_left = r.RooRealVar("n_Dplus_left","n_Dplus_left",power_Dplus_l) sig_Dplus_left = r.RooCBShape("signal Dplus left","signal Dplus left", m, mean_Dplus, sigma_Dplus, al_Dplus_left, pwr_Dplus_left) #Creating D+ signal PDF right tail al_Dplus_right = r.RooRealVar("alpha_Dplus_right","alpha_Dplus_right",alpha_Dplus_r) pwr_Dplus_right = r.RooRealVar("n_Dplus_right","n_Dplus_right",power_Dplus_r) sig_Dplus_right = r.RooCBShape("signal Dplus right","signal Dplus right", m, mean_Dplus, sigma_Dplus, al_Dplus_right, pwr_Dplus_right) #Adding the two CBs with their relative fraction sig_Dplus = r.RooAddPdf("signal Dplus","Dplus mass peak",r.RooArgList(sig_Dplus_left, sig_Dplus_right),r.RooArgList(sig_frc_Dplus)) #Creating Ds+ signal PDF mean_Ds= r.RooRealVar("mean_Ds","mean_Ds",1.97,1.93,2.01) sigma_Ds = r.RooRealVar("sigma_Ds","sigma_Ds",0.020,0.,0.2) sig_frc_Ds =r.RooRealVar("Ds Double crystalball fraction","Ds Double crystalball fraction",frac_Ds) al_Ds_left = r.RooRealVar("alpha_Ds_left","alpha_Ds_left",alpha_Ds_l) pwr_Ds_left = r.RooRealVar("n_Ds_left","n_Ds_left",power_Ds_l) sig_Ds_left = r.RooCBShape("signal Ds left","signal Ds left", m, mean_Ds, sigma_Ds, al_Ds_left, pwr_Ds_left) #Creating D+ signal PDF right tail al_Ds_right = r.RooRealVar("alpha_Ds_right","alpha_Ds_right",alpha_Ds_r) pwr_Ds_right = r.RooRealVar("n_Ds_right","n_Ds_right",power_Ds_r) sig_Ds_right = r.RooCBShape("signal Ds right","signal Ds right", m, mean_Ds, sigma_Ds, al_Ds_right, pwr_Ds_right) #Adding the two CBs with their relative fraction sig_Ds = r.RooAddPdf("sig_Ds","Ds mass peak",r.RooArgList(sig_Ds_left, sig_Ds_right),r.RooArgList(sig_frc_Ds)) #Model the background coef0 = r.RooRealVar("c0","coefficient #0",1.0,-10.,10) coef1 = r.RooRealVar("c1","coefficient #1",0.1,-10.,10) coef2 = r.RooRealVar("c2","coefficient #2",-0.1,-10.,10) bkg = r.RooChebychev("bkg","background p.d.f.",m,r.RooArgList(coef0,coef1,coef2)) #add it altogether #Add 2 sources of signal if mother_ID[mother_index_fit]=='both': NDs = r.RooRealVar("nDs","nDs",100.,0.,40000.) NDplus = r.RooRealVar("nDplus","nDplus",100.,0.,40000.) nbkg = r.RooRealVar("nbkg","nbkg",100.,0.,40000.) model = r.RooAddPdf("model","Dplus and Ds mass peaks",r.RooArgList(sig_Ds, sig_Dplus, bkg),r.RooArgList(NDs, NDplus,nbkg)) if mother_ID[mother_index_fit]=='Ds': NDs = r.RooRealVar("nDs","nDs",100.,0.,40000.) nbkg = r.RooRealVar("nbkg","nbkg",100.,0.,40000.) model = r.RooAddPdf("model","Ds mass peak",r.RooArgList(sig_Ds, bkg),r.RooArgList(NDs,nbkg)) if mother_ID[mother_index_fit]=='Dplus': NDplus = r.RooRealVar("nDplus","nDplus",100.,0.,40000.) nbkg = r.RooRealVar("nbkg","nbkg",100.,0.,40000.) model = r.RooAddPdf("model","Dplus mass peak",r.RooArgList(sig_Dplus, bkg),r.RooArgList(NDplus,nbkg)) fitr = model.fitTo(data_set,r.RooFit.Extended(),r.RooFit.Save()) xframe = m.frame(r.RooFit.Title("Fit to "+l_flv[l_index]+" data")) bkg_component = r.RooArgSet(bkg) sig_Dplus_component = r.RooArgSet(sig_Dplus) sig_Ds_component = r.RooArgSet(sig_Ds) #sig_D_component = r.RooArgSet(sig_Dplus_component,sig_Ds_component) data_set.plotOn(xframe) model.plotOn(xframe) hpull=xframe.pullHist() n_param = fitr.floatParsFinal().getSize() reduced_chi_square = xframe.chiSquare(n_param) model.plotOn(xframe,r.RooFit.Components(bkg_component),r.RooFit.LineStyle(2)) if mother_ID[mother_index_fit]=='Dplus': model.plotOn(xframe, r.RooFit.Components(sig_Dplus_component), r.RooFit.LineColor(2),r.RooFit.LineStyle(2) ) if mother_ID[mother_index_fit]=='Ds': model.plotOn(xframe, r.RooFit.Components(sig_Ds_component), r.RooFit.LineColor(2),r.RooFit.LineStyle(2) ) if mother_ID[mother_index_fit]=='both': model.plotOn(xframe, r.RooFit.Components(sig_Ds_component), r.RooFit.LineColor(2),r.RooFit.LineStyle(2) ) model.plotOn(xframe, r.RooFit.Components(sig_Dplus_component), r.RooFit.LineColor(2),r.RooFit.LineStyle(2) ) model.paramOn(xframe, r.RooFit.Layout(0.69,0.99,0.92), r.RooFit.Format("NEU", r.RooFit.AutoPrecision(1))) xframe.getAttText().SetTextSize(0.04) xframe2 = m.frame(r.RooFit.Title("Pulls")) xframe2.addPlotable(hpull,"P") c = r.TCanvas("Fit {0}".format(iteration),"Fit {0}".format(iteration),900,600) pad1 = r.TPad("pad1", "pad1", 0, 0.35, 1, 1.0) pad1.SetBottomMargin(0) pad1.Draw() c.cd() pad2 = r.TPad("pad2", "pad2", 0, 0.05, 1, 0.35) pad2.SetTopMargin(0) pad2.SetBottomMargin(0.2) pad2.SetGridx() pad2.SetGridy() pad2.Draw() pad1.cd() xframe.Draw() #xframe.GetYaxis().SetLabelOffset(-0.01) pad2.cd() xframe2.Draw() xframe2.GetXaxis().SetTitleSize(0.09) xframe2.GetXaxis().SetLabelSize(0.1) xframe2.GetYaxis().SetLabelSize(0.05) print("chi2 {0}".format(reduced_chi_square)) c.SaveAs(l_flv[l_index]+"/fits/"+mother_ID[mother_index_fit]+"/{0}/fit_{1}.png".format(iteration,iteration)) xframe_sw = m.frame(r.RooFit.Title("Weighted "+l_flv[l_index]+" data")) if mother_ID[mother_index_fit]=='both': splot = r.RooStats.SPlot( 'sData','sData', data_set, model, r.RooArgList(NDs, NDplus, nbkg) ) Ds_sw=splot.GetSWeightVars()[0] Dplus_sw=splot.GetSWeightVars()[1] bkg_sw=splot.GetSWeightVars()[2] argset = r.RooArgSet(m, Ds_sw) dataset_Ds_w = r.RooDataSet("Ds weighted dataset","Ds weighted dataset",data_set, argset,"","nDs_sw") argset = r.RooArgSet(m, Dplus_sw) dataset_Dplus_w = r.RooDataSet("Dplus weighted dataset","Dplus weighted dataset",data_set, argset,"","nDplus_sw") dataset_Ds_w.plotOn(xframe_sw, r.RooFit.MarkerSize(0.5), r.RooFit.MarkerColor(r.kBlue)) dataset_Dplus_w.plotOn(xframe_sw, r.RooFit.MarkerSize(0.5), r.RooFit.MarkerColor(r.kRed)) Ds_sw.setRange(-1.2,1.5) frame_Ds_sw = Ds_sw.frame(r.RooFit.Title("signal Ds sWeights")) data_set.plotOn(frame_Ds_sw, r.RooFit.MarkerSize(0.05)) Dplus_sw.setRange(-1.2,1.5) frame_Dplus_sw = Dplus_sw.frame(r.RooFit.Title("signal Dplus sWeights")) data_set.plotOn(frame_Dplus_sw, r.RooFit.MarkerSize(0.05)) bkg_sw.setRange(-.6,1.2) frame_bkg_sw = bkg_sw.frame(r.RooFit.Title("background sWeights")) data_set.plotOn(frame_bkg_sw, r.RooFit.MarkerSize(0.05)) sframe = m.frame(r.RooFit.Title("m vs sWeights")) data_set.plotOnXY(sframe, r.RooFit.YVar(Ds_sw), r.RooFit.MarkerColor(r.kGreen), r.RooFit.Name("Ds_sig")) data_set.plotOnXY(sframe, r.RooFit.YVar(Dplus_sw), r.RooFit.MarkerColor(r.kBlue), r.RooFit.Name("Dplus_sig")) data_set.plotOnXY(sframe, r.RooFit.YVar(bkg_sw), r.RooFit.MarkerColor(r.kRed), r.RooFit.Name("bkg")) legend = r.TLegend(0.89, 0.89, 0.5, 0.7) legend.AddEntry(sframe.findObject('Ds_sig'), 'Signal Ds sWeights', 'p') legend.AddEntry(sframe.findObject('Dplus_sig'), 'Signal Dplus sWeights', 'p') legend.AddEntry(sframe.findObject('bkg'), 'Background sWeights', 'p') # c = r.TCanvas("sWeighted mass dist {0}".format(iteration),"sWeighted mass dist {0}".format(iteration),900,600) xframe_sw.Draw() c.SaveAs(l_flv[l_index]+"/fits/"+mother_ID[mother_index_fit]+"/{0}".format(iteration)+"/sweighted_mass_plot_fit{0}.png".format(iteration)) c = r.TCanvas("sWeights {0}".format(iteration),"sWeights {0}".format(iteration),900,900) c.Divide(2,2) c.cd(1) frame_Ds_sw.Draw() c.cd(2) frame_Dplus_sw.Draw() c.cd(3) frame_bkg_sw.Draw() c.cd(4) sframe.Draw() legend.Draw() c.SaveAs(l_flv[l_index]+"/fits/"+mother_ID[mother_index_fit]+"/{0}".format(iteration)+"/sweights_{0}.png".format(iteration)) nbkg_w = splot.GetYieldFromSWeight("nbkg_sw") nDs_w = splot.GetYieldFromSWeight("nDs_sw") nDplus_w = splot.GetYieldFromSWeight("nDplus_sw") raw_sig=nDs_w+nDplus_w if mother_ID[mother_index_fit]=='Ds': splot = r.RooStats.SPlot( 'sData','sData', data_set, model, r.RooArgList(NDs, nbkg) ) Ds_sw=splot.GetSWeightVars()[0] bkg_sw=splot.GetSWeightVars()[1] argset = r.RooArgSet(m, Ds_sw) dataset_Ds_w = r.RooDataSet("Ds weighted dataset","Ds weighted dataset",data_set, argset,"","nDs_sw") dataset_Ds_w.plotOn(xframe_sw, r.RooFit.MarkerSize(0.5), r.RooFit.MarkerColor(r.kBlue)) Ds_sw.setRange(-3.,1.4) frame_Ds_sw = Ds_sw.frame(r.RooFit.Title("signal Ds sWeights")) data_set.plotOn(frame_Ds_sw, r.RooFit.MarkerSize(0.05)) bkg_sw.setRange(-.5,5.) frame_bkg_sw = bkg_sw.frame(r.RooFit.Title("background sWeights")) data_set.plotOn(frame_bkg_sw, r.RooFit.MarkerSize(0.05)) sframe = m.frame(r.RooFit.Title("m vs sWeights")) data_set.plotOnXY(sframe, r.RooFit.YVar(Ds_sw), r.RooFit.MarkerColor(r.kGreen), r.RooFit.Name("Ds_sig")) data_set.plotOnXY(sframe, r.RooFit.YVar(bkg_sw), r.RooFit.MarkerColor(r.kRed), r.RooFit.Name("bkg")) legend = r.TLegend(0.89, 0.89, 0.75, 0.8) legend.AddEntry(sframe.findObject('Ds_sig'), 'Signal Ds sWeights', 'p') legend.AddEntry(sframe.findObject('bkg'), 'Background sWeights', 'p') # c = r.TCanvas("sWeighted mass dist {0}".format(iteration),"sWeighted mass dist {0}".format(iteration),900,600) xframe_sw.Draw() c.SaveAs(l_flv[l_index]+"/fits/"+mother_ID[mother_index_fit]+"/{0}".format(iteration)+"/sweighted_mass_plot_fit{0}.png".format(iteration)) c = r.TCanvas("sWeights {0}".format(iteration),"sWeights {0}".format(iteration),900,600) c.cd() pad1 = r.TPad("pad1", "pad1", 0, 0.5, 0.5, 1.0) pad1.Draw() c.cd() pad2 = r.TPad("pad2", "pad2", 0.5, 0.5, 1.0, 1.0) pad2.Draw() c.cd() pad3 = r.TPad("pad3", "pad3", 0, 0.05, 1, 0.50) pad3.SetGridx() pad3.Draw() pad1.cd() frame_Ds_sw.Draw() pad2.cd() frame_bkg_sw.Draw() pad3.cd() sframe.Draw() legend.Draw() c.SaveAs(l_flv[l_index]+"/fits/"+mother_ID[mother_index_fit]+"/{0}".format(iteration)+"/sweights_{0}.png".format(iteration)) nbkg_w = splot.GetYieldFromSWeight("nbkg_sw") nDs_w = splot.GetYieldFromSWeight("nDs_sw") raw_sig=nDs_w if mother_ID[mother_index_fit]=='Dplus': splot = r.RooStats.SPlot( 'sData','sData', data_set, model, r.RooArgList(NDplus, nbkg) ) Dplus_sw=splot.GetSWeightVars()[0] bkg_sw=splot.GetSWeightVars()[1] argset = r.RooArgSet(m, Dplus_sw) dataset_Dplus_w = r.RooDataSet("Dplus weighted dataset","Dplus weighted dataset",data_set, argset,"","nDplus_sw") #model.plotOn(xframe_sw, r.RooFit.Components(sig_Dplus_component),r.RooFit.LineColor(2),r.RooFit.LineStyle(2) ) dataset_Dplus_w.plotOn(xframe_sw, r.RooFit.MarkerSize(0.5), r.RooFit.MarkerColor(r.kRed)) Dplus_sw.setRange(-1.2,1.8) frame_Dplus_sw = Dplus_sw.frame(r.RooFit.Title("signal Dplus sWeights")) data_set.plotOn(frame_Dplus_sw, r.RooFit.MarkerSize(0.05)) bkg_sw.setRange(-1.2,2.1) frame_bkg_sw = bkg_sw.frame(r.RooFit.Title("background sWeights")) data_set.plotOn(frame_bkg_sw, r.RooFit.MarkerSize(0.05)) sframe = m.frame(r.RooFit.Title("m vs sWeights")) data_set.plotOnXY(sframe, r.RooFit.YVar(Dplus_sw), r.RooFit.MarkerColor(r.kBlue), r.RooFit.Name("Dplus_sig")) data_set.plotOnXY(sframe, r.RooFit.YVar(bkg_sw), r.RooFit.MarkerColor(r.kRed), r.RooFit.Name("bkg")) legend = r.TLegend(0.89, 0.89, 0.75, 0.8) legend.AddEntry(sframe.findObject('Dplus_sig'), 'Signal Dplus sWeights', 'p') legend.AddEntry(sframe.findObject('bkg'), 'Background sWeights', 'p') # c = r.TCanvas("sWeighted mass dist {0}".format(iteration),"sWeighted mass dist {0}".format(iteration),900,600) xframe_sw.Draw() c.SaveAs(l_flv[l_index]+"/fits/"+mother_ID[mother_index_fit]+"/{0}".format(iteration)+"/sweighted_mass_plot_fit{0}.png".format(iteration)) c = r.TCanvas("sWeights {0}".format(iteration),"sWeights {0}".format(iteration),900,600) c.cd() pad1 = r.TPad("pad1", "pad1", 0, 0.5, 0.5, 1.0) pad1.Draw() c.cd() pad2 = r.TPad("pad2", "pad2", 0.5, 0.5, 1.0, 1.0) pad2.Draw() c.cd() pad3 = r.TPad("pad3", "pad3", 0, 0.05, 1, 0.50) pad3.SetGridx() pad3.Draw() pad1.cd() frame_Dplus_sw.Draw() pad2.cd() frame_bkg_sw.Draw() pad3.cd() sframe.Draw() legend.Draw() c.SaveAs(l_flv[l_index]+"/fits/"+mother_ID[mother_index_fit]+"/{0}".format(iteration)+"/sweights_{0}.png".format(iteration)) nbkg_w = splot.GetYieldFromSWeight("nbkg_sw") nDplus_w = splot.GetYieldFromSWeight("nDplus_sw") raw_sig=nDplus_w fom=(Nsig_from_MC)/np.sqrt(nbkg_w+Nsig_from_MC) eff_corrected_sig = (raw_sig)/(sig_sel_eff) fit_results={ '# XGBoost test {0} BDT output'.format(test): '\n', '# background': nbkg_w, #'err bkg': err_bkg, '# fitted signal': raw_sig, #'error sig': raw_sig_err, '# of signal in MC': Nsig_from_MC, 'signal selection efficiency': sig_sel_eff, 'S/sqrt(S+B)': fom, 'efficiency corrected yeld': eff_corrected_sig, 'chi 2': reduced_chi_square } return fit_results if __name__=='__main__': parser = argparse.ArgumentParser(description='Fit at different BDT cuts') parser.add_argument('-iteration', metavar='Fit iteration', type =int, help='Fit iteration') #parser.add_argument('-test', metavar='BDT test', type =int, help='Which BDT test you are using') parser.add_argument('-x_cut', metavar='BDT cut', type =float, help='BDT cut') parser.add_argument('-l_index', metavar='Lepton flavour', type =int, help='Which lepton flavour are you fitting?') parser.add_argument('-mother_index_fit', metavar='Decay mother', type =int, help='Which mass window are you fitting?') args = parser.parse_args() test=4 #test=args.test i = args.iteration l_index = args.l_index mother_index_fit = args.mother_index_fit x_cut=args.x_cut FIT_PATH=l_flv[l_index]+'/fits' BDT_PATH=l_flv[l_index]+'/BDTs/test_'+str(test) if not os.path.exists(FIT_PATH+'/'+mother_ID[mother_index_fit]+'/{0}/'.format(i)): os.mkdir(FIT_PATH+'/'+mother_ID[mother_index_fit]+'/{0}/'.format(i)) fit_results = select_and_fit(mother_index_fit, l_index, x_cut, i, test) with open(FIT_PATH+'/'+mother_ID[mother_index_fit]+'/fit_results_{0}.txt'.format(i), 'w') as g: print('Saving fit results...') for key, value in fit_results.items(): g.write('%s:%s\n' % (key, value)) g.close()