import ROOT as r import root_numpy as rn import pickle import numpy as np import matplotlib.pyplot as plt import array as array from xgboost import XGBClassifier from tools.data_processing import * import os l_flv=['e','mu'] mother_ID=["Ds","Dplus","both"] def MC_fit(mc_mass, mother_index, l_index): FIT_PATH=l_flv[l_index]+'/fits' print("Fitting MC "+mother_ID[mother_index]+" signal") if l_flv[l_index]=='mu': if mother_ID[mother_index]=='Dplus': lower_mass_cut=mc_mass.min() upper_mass_cut=mc_mass.max() if mother_ID[mother_index]=='Ds': lower_mass_cut=mc_mass.min() upper_mass_cut=mc_mass.max() if l_flv[l_index]=='e': if mother_ID[mother_index]=='Dplus': lower_mass_cut=mc_mass.min() upper_mass_cut=mc_mass.max() if mother_ID[mother_index]=='Ds': lower_mass_cut=mc_mass.min() upper_mass_cut=mc_mass.max() mc_fit=mc_mass/1000. lower_mass_cut/=1000. upper_mass_cut/=1000. #create a rootfile from which will do the unbinned fit data_array=np.array([mc_fit[i] for i in range(len(mc_fit))], dtype=[('MC_D_m', np.float32)]) rn.array2root(data_array, filename='/disk/lhcb_data/davide/Rphipi/MC/for_fit/'+l_flv[l_index]+l_flv[l_index]+'/MC_'+mother_ID[mother_index]+'_'+l_flv[l_index]+l_flv[l_index]+'_mass.root', treename='D_M', mode='recreate', ) f=r.TFile('/disk/lhcb_data/davide/Rphipi/MC/for_fit/'+l_flv[l_index]+l_flv[l_index]+'/MC_'+mother_ID[mother_index]+'_'+l_flv[l_index]+l_flv[l_index]+'_mass.root') tree=f.Get("D_M") #some tricks to pass Roofit a dataset in python mass = np.array([0],dtype=np.float32) branch = tree.GetBranch("MC_D_m") branch.SetAddress(mass) num_entries=tree.GetEntries() m = r.RooRealVar("m","m",lower_mass_cut,upper_mass_cut) l = r.RooArgSet(m) data_hist = r.RooDataSet("data histogram", "data", l) for i in range(num_entries): tree.GetEvent(i) r.RooAbsRealLValue.__assign__(m, mass[0]) data_hist.add(l, 1.0) #Creating the signal PDF as two crystall balls if l_flv[l_index]=='mu': if mother_ID[mother_index]=='Dplus': mean_d_startval=1.85 mean_d_up=1.95 mean_d_down=1.75 if mother_ID[mother_index]=='Ds': mean_d_startval=1.95 mean_d_up=2.05 mean_d_down=1.85 if l_flv[l_index]=='e': if mother_ID[mother_index]=='Dplus': mean_d_startval=1.85 mean_d_up=2.30 mean_d_down=1.70 if mother_ID[mother_index]=='Ds': mean_d_startval=1.95 mean_d_up=2.30 mean_d_down=1.70 #Creating D signal PDF left tail mean_d= r.RooRealVar("mean_D","mean_D",mean_d_startval,mean_d_down,mean_d_up) sigma_d = r.RooRealVar("sigma_D","sigma_D",0.020,0.,0.2) alpha_l = r.RooRealVar("alpha_D_l","alpha_D_l",-0.020,-10.,10.) n_l = r.RooRealVar("n_D_l","n_D_l",0.020,-50.,50.) sig_l = r.RooCBShape("signal_D_l","signal_D_l", m, mean_d, sigma_d, alpha_l, n_l) #Creating D signal PDF right tail alpha_r = r.RooRealVar("alpha_D_r","alpha_D_r",0.020,-10.,10.) n_r = r.RooRealVar("n_D_r","n_D_r",0.020,-50.,50.) sig_r = r.RooCBShape("signal_D_r","signal_D_r", m, mean_d, sigma_d, alpha_r, n_r) #Add 2 sources of signal # first add the sigs together frac = r.RooRealVar("# of Ds sig events","# of Ds sig events",0.5,0.,1.) model = r.RooAddPdf("model","two mass peaks",r.RooArgList(sig_l,sig_r),r.RooArgList(frac)) fitr = model.fitTo(data_hist,r.RooFit.Save()) xframe = m.frame(r.RooFit.Title("Fit to "+l_flv[l_index]+" MC signal")) data_hist.plotOn(xframe) model.plotOn(xframe) c = r.TCanvas() xframe.Draw() n_param = fitr.floatParsFinal().getSize() reduced_chi_square = xframe.chiSquare(n_param) CB_fraction=frac.getVal() Nsig_from_MC=mc_mass.shape[0] a_l = alpha_l.getVal() pwr_l= n_l.getVal() a_r= alpha_r.getVal() pwr_r=n_r.getVal() print("chi2 {0}".format(reduced_chi_square)) c.SaveAs(l_flv[l_index]+"/fits/MC_"+mother_ID[mother_index]+"_fit.png") MC_fit_results={ '# Fitting the MC signal':'\n', 'CB fraction': CB_fraction, '# of signal in MC': Nsig_from_MC, 'alpha_'+mother_ID[mother_index]+'_l': a_l, 'n_'+mother_ID[mother_index]+'_l': pwr_l, 'alpha_'+mother_ID[mother_index]+'_r': a_r, 'n_'+mother_ID[mother_index]+'_r': pwr_r, } print('Saving MC fit results...') with open(FIT_PATH+'/MC_'+mother_ID[mother_index]+'_fit_results.txt', 'w') as f: for key, value in MC_fit_results.items(): f.write('%s:%s\n' % (key, value)) with open(FIT_PATH+'/MC_'+mother_ID[mother_index]+'_CB_params.pickle','w') as handle: pickle.dump(MC_fit_results, handle) print('Saved MC fit results \n Proceeding to fit script') return MC_fit_results, a_l, pwr_l, a_r, pwr_r, CB_fraction