#!/usr/bin/env python # coding: utf-8 import ROOT #import numpy as np from array import array import os import sys # Version with copy of the original tree: #take path given as argument filelocation=sys.argv[1] #define the outputs directory if (("/Collision11/" in filelocation)&("/00167485/0000/" in filelocation)): sample="2011Down" path="/eos/lhcb/grid/prod/lhcb/LHCb/Collision11/PIDCALIB.ROOT/00167485/0000/" elif (("/Collision11/" in filelocation)&("/00167487/0000/" in filelocation)): sample="2011Up" path="/eos/lhcb/grid/prod/lhcb/LHCb/Collision11/PIDCALIB.ROOT/00167487/0000/" elif (("/Collision12/" in filelocation)&("/00165235/0000/" in filelocation)): sample="2012Down" path="/eos/lhcb/grid/prod/lhcb/LHCb/Collision12/PIDCALIB.ROOT/00165235/0000/" elif (("/Collision12/" in filelocation)&("/00165237/0000/" in filelocation)): sample="2012Up" path="/eos/lhcb/grid/prod/lhcb/LHCb/Collision12/PIDCALIB.ROOT/00165237/0000/" #if (filelocation.replace( "base_ntuples" ,"weighted_ntuples") in newfiles): continue #if sWeighting has been already done for a file -> don't (re)do it print(sample, path) outputfile_name = filelocation.replace(path,f"/eos/lhcb/user/t/tmiralle/sFitted_tuples/{sample}/weighted_ntuples/") out=ROOT.TFile(outputfile_name, "RECREATE") out.Close() intermediatefile_name = filelocation.replace(path,f"/eos/lhcb/user/t/tmiralle/sFitted_tuples/{sample}/intermediate_ntuples/") intermediate=ROOT.TFile(intermediatefile_name, "RECREATE") intermediate.Close() finalfile_name = filelocation.replace(path,f"/eos/lhcb/user/t/tmiralle/sFitted_tuples/{sample}/final_ntuples/") final=ROOT.TFile(finalfile_name, "RECREATE") final.Close() for part in ['Pi','K']: ch = ROOT.TChain(f'DSt_{part}Tuple/DecayTree') ch.Add(filelocation) print( "Open " + filelocation) # c1=ROOT.TCanvas() ch.Draw('Dst_M') c1.Draw() #take the needed variables lvars=ROOT.RooArgSet() mDstar = ROOT.RooRealVar("Dst_M","Dst_M",1940.,2080.) lvars.add(mDstar) mD0 = ROOT.RooRealVar("Dz_M","Dz_M",1800.,1920.) lvars.add(mD0) #make root dataset with needed columns data = ROOT.RooDataSet("data","data",ch,lvars) #mDst-mD0 deltaM_func = ROOT.RooFormulaVar( "deltaM" , "@0-@1" , ROOT.RooArgList( mDstar , mD0 )) deltaM=data.addColumn(deltaM_func) deltaM.setRange(141.,152.) #2 gaussian model to fit mD0 mean = ROOT.RooRealVar("mean","mean",1865,1840.,1900.) sigma = ROOT.RooRealVar("sigma","sigma",6.7,4.,20.) gauss = ROOT.RooGaussian("gauss","gauss",mD0,mean,sigma) sigma2 = ROOT.RooRealVar("sigma2","sigma2",11.,0.,100.) gauss2 = ROOT.RooGaussian("gauss2","gauss2",mD0,mean,sigma2) #Chebychev model to fit bkg part a0 = ROOT.RooRealVar("a0","a0",0.,-1.,10.)#allow negative slop bkg=ROOT.RooChebychev('bkg','bkg',mD0,ROOT.RooArgList(a0)) nsig=ROOT.RooRealVar("nSig","nSig",80000,0.,100000000.) #frac of the two gaussian for the signal fracG=ROOT.RooRealVar("fracG","fracG",0.5,0.,1.) signal = ROOT.RooAddPdf( "signal","signal",gauss,gauss2,fracG) nbkg=ROOT.RooRealVar("nBkg","nBkg",19000,0,1000000000.) #total model model=ROOT.RooAddPdf("model","model",ROOT.RooArgList(signal,bkg),ROOT.RooArgList(nsig,nbkg)) #plots of the data before deltaM cut c22=ROOT.TCanvas() deltaMframe = mDstar.frame() data.plotOn(deltaMframe) deltaMframe.Draw() c22.Draw('Dst_M') #cut on delta m #data_red = data.reduce("(deltaM>143)&(deltaM<148)")#removed cut on delta M because it seem not used in Run2 pidcalib.root data_red = data #plots of the data after deltaM cut c2=ROOT.TCanvas() deltaMframe = mDstar.frame() data_red.plotOn(deltaMframe) deltaMframe.Draw() c2.Draw('Dst_M') c3=ROOT.TCanvas() deltaMframe = mD0.frame() data_red.plotOn(deltaMframe) deltaMframe.Draw() c3.Draw('Dz_M') #fit of the data and plot of the fit, with all free parameter model.fitTo(data_red,ROOT.RooFit.Extended(),ROOT.RooFit.NumCPU(7)) deltaMframe = mD0.frame() data_red.plotOn(deltaMframe) model.plotOn(deltaMframe) deltaMframe.Draw() c3.Draw() # sPlot, parameters, except the yields,fixed via previous fit mean.setConstant() sigma.setConstant() sigma2.setConstant() a0.setConstant() fracG.setConstant() sData = ROOT.RooStats.SPlot("sData", "sPlot", data_red, model, ROOT.RooArgList( nsig, nbkg)) print("Signal yield w.r.t. sWeights: {}".format(sData.GetYieldFromSWeight("nSig"))) print("Bkg yield w.r.t. sWeights: {}".format(sData.GetYieldFromSWeight("nBkg"))) print("Total number of events for the sFit: {}".format(data_red.sumEntries())) c5=ROOT.TCanvas() model.fitTo(data_red,ROOT.RooFit.Extended(),ROOT.RooFit.NumCPU(7)) deltaMframe = mD0.frame() data_red.plotOn(deltaMframe) model.plotOn(deltaMframe) #model.plotOn(deltaMframe,Components={signal},LineStyle="--")#don't work with my version #model.plotOn(deltaMframe,Components={bkg},LineStyle=":") deltaMframe.Draw() c5.Draw() if sData.GetYieldFromSWeight("nSig") < 0: continue # c8=ROOT.TCanvas() nSig_sw=sData.GetSWeightVars().find('nSig_sw')#nSig_sw->probe_sWeight (5 times here) nSig_sw.SetTitle('Signal sWeight') nSig_sw.setRange(-1.5, 1.5) frame_sigsw = nSig_sw.frame(ROOT.RooFit.Title('Signal sWeights')) data_red.plotOn(frame_sigsw) frame_sigsw.Draw() c8.Draw() # ROOT.RooDataSet.setDefaultStorageType(ROOT.RooAbsData.Tree) out=ROOT.TFile(outputfile_name, "UPDATE") out.mkdir(f"DSt_{part}Tuple/") out.cd(f"DSt_{part}Tuple/") sTree=ROOT.RooStats.GetAsTTree(f'sTree', 'sTree', sData.GetSDataSet() ) sTree.Write() out.Close() ROOT.RooDataSet.setDefaultStorageType(ROOT.RooAbsData.Vector) #copy original tree in intermediate file original=ROOT.TFile(filelocation,"READ") inTree=original.Get(f'DSt_{part}Tuple/DecayTree') nb_events=inTree.GetEntries() print(inTree.GetEntries('(Dst_M>1940)&(Dst_M<2080)&(Dz_M>1800)&(Dz_M<1920)')) intermediate=ROOT.TFile(intermediatefile_name, "UPDATE") intermediate.mkdir(f"DSt_{part}Tuple/") intermediate.cd(f"DSt_{part}Tuple/") copTree=inTree.CopyTree("(Dst_M>1940)&(Dst_M<2080)&(Dz_M>1800)&(Dz_M<1920)") copTree.Write() original.Close() intermediate.Close() #put sWeight branch on the Tree into the intermediate file out=ROOT.TFile(outputfile_name, "READ") sTree=out.Get(f'DSt_{part}Tuple/sTree') intermediate=ROOT.TFile(intermediatefile_name, "UPDATE") intermediate.cd(f"DSt_{part}Tuple/") copTree=intermediate.Get(f'DSt_{part}Tuple/DecayTree') nb_events=sTree.GetEntries() probe_sWeight= array('d',[0]) newBranch =copTree.Branch('probe_sWeight', probe_sWeight, 'probe_sWeight/D')#newbranch to fill with sWeights for i in range (nb_events): sTree.GetEntry(i) probe_sWeight[0]=sTree.nSig_sw newBranch.Fill() copTree.Write() out.Close() intermediate.Close() #There is a duplication of tree's in intermediate file, so: #copy final tree with sw in final file intermediate=ROOT.TFile(intermediatefile_name,"READ") FinTree=intermediate.Get(f'DSt_{part}Tuple/DecayTree;2') final=ROOT.TFile(finalfile_name, "UPDATE") final.mkdir(f"DSt_{part}Tuple/") final.cd(f"DSt_{part}Tuple/") copTree2=FinTree.CopyTree("") copTree2.Write() intermediate.Close() final.Close() #loop to fill original not used tree in the final file other_subdir=['Lam0_PTuple','Jpsi_MuMTuple','Jpsi_MuPTuple','Jpsi_eeTuple','GetIntegratedLuminosity'] for name in other_subdir: original=ROOT.TFile(filelocation,"READ") if (name=='GetIntegratedLuminosity'): inTree=original.Get(f'{name}/LumiTuple') final=ROOT.TFile(finalfile_name, "UPDATE") final.mkdir(f"{name}/") final.cd(f"{name}/") copTree=inTree.CopyTree("") copTree.Write() original.Close() final.Close() else: inTree=original.Get(f'{name}/DecayTree') final=ROOT.TFile(finalfile_name, "UPDATE") final.mkdir(f"{name}/") final.cd(f"{name}/") copTree=inTree.CopyTree("") copTree.Write() original.Close() final.Close() os.remove(intermediatefile_name) os.remove(outputfile_name) print('OK')