import ROOT import rootnotes import rootprint import utils from ProgressBar import ProgressBar import math from ROOT import RooRealVar, RooFormulaVar, RooVoigtian, RooChebychev, RooArgList, \ RooArgSet, RooAddPdf, RooDataSet, RooCategory, RooSimultaneous, \ RooBreitWigner, RooCBShape, RooFFTConvPdf ROOT.gROOT.LoadMacro("Loader.C+") # print open("Loader.C").read() tree = ROOT.TChain("T") tree.Add("/afs/cern.ch/user/d/demattia/work/public/TagAndProbe_ZMuMu.root") mass = RooRealVar("mass", "mass", 60, 120) # Construct signal pdf mean = RooRealVar("mean", "mean", 90, 80, 100) width = RooRealVar("width", "width", 2.4952, 1, 3) width.setConstant(ROOT.kTRUE) sigma = RooRealVar("sigma", "sigma", 1.2, 0.2, 2) signal = RooVoigtian("signal", "signal", mass, mean, width, sigma) # Construct background pdf a0 = RooRealVar("a0","a0",0.1,-1,1) a1 = RooRealVar("a1","a1",0.004,-1,1) background = RooChebychev("background","background",mass,RooArgList(a0,a1)) # Construct composite pdf nsig = RooRealVar("nsig", "nsig", 2000, 0, 100000) nbkg = RooRealVar("nbkg", "nbkg", 100, 0, 1000) model = RooAddPdf("model", "model", RooArgList(signal, background), RooArgList(nsig, nbkg)) # alpha = RooRealVar("alpha", "alpha", 80, 60, 100) # CBn = RooRealVar("n", "n", 1, -10, 10) # BW = RooBreitWigner("BW","Breit Wigner theory",mass,mean,width) # cryBall = RooCBShape("cryBall","Crystal Ball resolution model", mass, mean, sigma, alpha, CBn) # model = RooAddPdf("model", "model", RooArgList(cryBall, background), RooArgList(nsig, nbkg)) import bisect ptBins = [20, 40, 70] # Returns the bins corresponding to the given pt values def findBins(pt1, pt2): return (bisect.bisect_right(ptBins, pt1)-1, bisect.bisect_right(ptBins, pt2)-1) # This returns the position in an single dimension array given the bins def findPosition(bin1, bin2): return bin1+(len(ptBins)-1)*bin2 datasetAllMap = {} for ptBin1 in range(0, len(ptBins)): for ptBin2 in range(0, len(ptBins)): datasetAllMap[(ptBin1, ptBin2)] = RooDataSet("datasetAll_"+str(ptBin1)+"_"+str(ptBin2), "datasetAll_"+str(ptBin1)+"_"+str(ptBin2), RooArgSet(mass)) hAllMap = {} for ptBin1 in range(0, len(ptBins)): for ptBin2 in range(0, len(ptBins)): hAllMap[(ptBin1,ptBin2)] = ROOT.TH1F("hAll"+"_"+str(ptBin1)+"_"+str(ptBin2), "All events passing old trigger "+str(ptBin1)+"_"+str(ptBin2), 100, 60, 120) muMass = 0.1057 def passSelection(track): if ( track.pt > 26 and abs(track.eta) < 2 and track.isolation/track.pt < 0.1 and track.trackerLayersWithMeasurement >= 6 and track.trackQuality and track.dxy < 30. and track.dz < 30. ): return True def passDileptonSelection(track1, track2): if track1.charge == track2.charge: return False if utils.deltaR(track1.phi, track1.eta, track2.phi, track2.eta) < 0.2: return False return True def computeMass(mu1, mu2): muon1 = ROOT.TLorentzVector() muon2 = ROOT.TLorentzVector() muon1.SetPtEtaPhiM(mu1.pt, mu1.eta, mu1.phi, muMass) muon2.SetPtEtaPhiM(mu2.pt, mu2.eta, mu2.phi, muMass) return (muon1+muon2).M() def fillSingleCandidate(track1, track2, histoMap, datasetMap): if not passDileptonSelection(track1, track2): return False mass.setVal(computeMass(track1,track2)) histoMap[findBins(track1.pt, track2.pt)].Fill(mass.getVal()) datasetMap[findBins(track1.pt, track2.pt)].add(RooArgSet(mass)) return True def fillCandidates(tracks, histoMap, datasetMap): if len(tracks) == 2: fillSingleCandidate(tracks[0], tracks[1], histoMap, datasetMap) elif len(tracks) > 2: for i in range(0, len(tracks)): for j in range(i+1, len(tracks)): fillSingleCandidate(tracks[i], tracks[j], histoMap, datasetMap) %%rootprint allCandidates = 0 passCandidates = 0 processedEvents = 0 maxEvents = 10000 # p = ProgressBar(maxEvents) for event in tree: if processedEvents > maxEvents: break # p.animate(processedEvents) processedEvents += 1 fillCandidates(event.muons, hAllMap, datasetAllMap) for bin1 in range(len(ptBins)-1): for bin2 in range(len(ptBins)-1): print "candidates["+str(ptBins[bin1])+"-"+str(ptBins[bin1+1])+\ ", "+str(ptBins[bin2])+"-"+str(ptBins[bin2+1])+"] =",\ hAllMap[bin1,bin2].GetEntries() canvas = rootnotes.canvas("AllCanvas", (800, 800)) canvas.Divide(len(ptBins)-1,len(ptBins)-1) for ptBin1 in range(0, len(ptBins)-1): for ptBin2 in range(0, len(ptBins)-1): canvas.cd(findPosition(ptBin1, ptBin2)+1) hAllMap[(ptBin1, ptBin2)].Draw() canvas #%%rootprint # Perform a fit of model to data # Use this for minos (better error estimate, but takes longer) # fr = model.fitTo(combData, ROOT.RooFit.Save(ROOT.kTRUE), ROOT.RooFit.Extended(ROOT.kTRUE), ROOT.RooFit.Minos(ROOT.kTRUE)) # fr = model.fitTo(combData, ROOT.RooFit.Save(ROOT.kTRUE), ROOT.RooFit.Extended(ROOT.kTRUE)) frMap = {} # Construct dataset in (x,sample) for ptBin1 in range(0, len(ptBins)-1): for ptBin2 in range(0, len(ptBins)-1): frMap[(ptBin1, ptBin2)] = model.fitTo(datasetAllMap[(ptBin1, ptBin2)], ROOT.RooFit.Save(ROOT.kTRUE), ROOT.RooFit.Extended(ROOT.kTRUE)) # Use this to run with Minos enabled. Minos allows for a more precise esimation of the parameter errors, # but the fit will take longer. # frMap[(ptBin1, ptBin2)] = model.fitTo(datasetAllMap[(ptBin1, ptBin2)], # ROOT.RooFit.Save(ROOT.kTRUE), # ROOT.RooFit.Extended(ROOT.kTRUE), # ROOT.RooFit.Minos(ROOT.kTRUE)) %%rootprint for ptBin1 in range(0, len(ptBins)-1): for ptBin2 in range(0, len(ptBins)-1): frMap[(ptBin1, ptBin2)].Print("v") def plotResults(ptBin1, ptBin2, Data, canvas2): frame2 = mass.frame(ROOT.RooFit.Bins(30),ROOT.RooFit.Title("All events")) # Plot all data tagged as physics sample Data.plotOn(frame2) model.plotOn(frame2) model.plotOn(frame2, ROOT.RooFit.Components("background"), ROOT.RooFit.LineStyle(ROOT.kDashed)) canvas2.cd(findPosition(ptBin1, ptBin2)+1) frame2.GetYaxis().SetTitleOffset(1.4) frame2.Draw() # Canvases for fit results canvas2 = ROOT.TCanvas("RooFitCanvas", "RooFitCanvas", 800, 800) canvas2.Divide(len(ptBins)-1,len(ptBins)-1) for ptBin1 in range(0, len(ptBins)-1): for ptBin2 in range(0, len(ptBins)-1): plotResults(ptBin1, ptBin2, datasetAllMap[(ptBin1,ptBin2)], canvas2) canvas2 #canvas2.Print("allFits.pdf")