Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2024-07-16 02:43:05

0001 import ROOT,itertools,math      
0002 import argparse
0003 from array import array          
0004 from DataFormats.FWLite import Events, Handle
0005 import subprocess
0006 ROOT.FWLiteEnabler.enable()
0007 
0008 
0009 class muon_trigger_analyzer(object):
0010      def __init__(self,prefix,thresholds=[0,3,5,10,15,20,30]):
0011           self.bunchfactor = 40000*2760.0/3564.0
0012           self.prefix=prefix
0013           self.thresholds = thresholds
0014           self.histoData={'effpt':{}, 'effeta':{},'geneta':{},'genphi':{},'genlxy':{},'efflxy':{},'effphi':{},'effptB':{},'effptO':{},'effptE':{},'rateeta':{}}
0015           self.histoData['genpt']         = ROOT.TH1D(prefix+"_genpt","genpt",50,0,100)
0016           self.histoData['genptB']        = ROOT.TH1D(prefix+"_genptB","genptB",50,0,100)
0017           self.histoData['genptO']        = ROOT.TH1D(prefix+"_genptO","genptO",50,0,100)
0018           self.histoData['genptE']        = ROOT.TH1D(prefix+"_genptE","genptE",50,0,100)
0019           self.histoData['rate']          = ROOT.TH1D(prefix+"_rate","rate",20,0,100)
0020           self.histoData['resolution']    = ROOT.TH2D(prefix+"_resolution","resolution",20,0,100,100,-2,2)
0021           self.histoData['rateBarrel']    = ROOT.TH1D(prefix+"_rateBarrel","rate",20,0,100)
0022           self.histoData['rateOverlap']   = ROOT.TH1D(prefix+"_rateOverlap","rate",20,0,100)
0023           self.histoData['rateEndcap']    = ROOT.TH1D(prefix+"_rateEndcap","rate",20,0,100)
0024           for t in thresholds:
0025                self.histoData['effpt'][t] = ROOT.TH1D(prefix+"_effpt_"+str(t),"effpt_"+str(t),50,0,100)
0026                self.histoData['effptB'][t] = ROOT.TH1D(prefix+"_effptB_"+str(t),"effpt_"+str(t),50,0,100)
0027                self.histoData['effptO'][t] = ROOT.TH1D(prefix+"_effptO_"+str(t),"effpt_"+str(t),50,0,100)
0028                self.histoData['effptE'][t] = ROOT.TH1D(prefix+"_effptE_"+str(t),"effpt_"+str(t),50,0,100)
0029                self.histoData['effeta'][t] = ROOT.TH1D(prefix+"_effeta_"+str(t),"effeta_"+str(t),48,-2.4,2.4)
0030                self.histoData['effphi'][t] = ROOT.TH1D(prefix+"_effphi_"+str(t),"effphi_"+str(t),32,-math.pi,math.pi)
0031                self.histoData['efflxy'][t] = ROOT.TH1D(prefix+"_efflxy_"+str(t),"efflxy_"+str(t),50,0,200)
0032                self.histoData['geneta'][t] = ROOT.TH1D(prefix+"_geneta_"+str(t),"effeta_"+str(t),48,-2.4,2.4)
0033                self.histoData['genphi'][t] = ROOT.TH1D(prefix+"_genphi_"+str(t),"genphi_"+str(t),32,-math.pi,math.pi)
0034                self.histoData['genlxy'][t] = ROOT.TH1D(prefix+"_genlxy_"+str(t),"genlxy_"+str(t),50,0,200)
0035                self.histoData['rateeta'][t] = ROOT.TH1D(prefix+"_rateeta_"+str(t),"rateeta_"+str(t),24,-2.4,2.4)
0036 
0037      def deltaPhi(self, p1, p2):
0038           '''Computes delta phi, handling periodic limit conditions.'''
0039           res = p1 - p2
0040           while res > math.pi:
0041                res -= 2*math.pi
0042           while res < -math.pi:
0043                res += 2*math.pi
0044           return res
0045 
0046      def deltaR(self, *args ):
0047           return math.sqrt( self.deltaR2(*args) )
0048 
0049      def deltaR2(self, e1, p1, e2, p2):
0050           de = e1 - e2
0051           dp = self.deltaPhi(p1, p2)
0052           return de*de + dp*dp
0053      def getLxy(self,muon):
0054           return abs(math.sqrt(muon.vx()*muon.vx()+muon.vy()*muon.vy()))
0055 
0056      def getDxy(self,muon):
0057           tanphi = math.tan(m.phi())
0058           x=(tanphi*tanphi*muon.vx()-muon.vy()*tanphi)/(1+tanphi*tanphi)
0059           y=muon.vy()+tanphi*(x-muon.vx())
0060           return abs(math.sqrt(x*x+y*y))
0061 
0062      def getEta1(self,muon):
0063           lxy = self.getLxy(muon)
0064           vz = muon.vz()
0065           theta1 = math.atan((700.-lxy)/(650.-vz))
0066           
0067           if (theta1 < 0):
0068                theta1 = math.pi+theta1
0069           eta1 = -math.log(math.tan(theta1/2.0))
0070           return eta1
0071 
0072      def getEta2(self,muon):
0073           lxy = self.getLxy(muon)
0074           vz = muon.vz()
0075           theta2 = math.pi-math.atan((700.-lxy)/(650.+vz))
0076           if theta2 > math.pi:
0077                theta2 = theta2-math.pi
0078           eta2 = -math.log(math.tan(theta2/2.0))
0079           return eta2                                     
0080 
0081      def getAcceptance(self,muon):
0082           eta1 = self.getEta1(muon)
0083           eta2 = self.getEta2(muon)
0084           if muon.eta() < eta1 and muon.eta() > eta2:
0085                return True #Muon is within barrel acceptance
0086           else:
0087                return False #Muon is outside of barrel 
0088 
0089      def getSt2Eta(self,muon):
0090           lxy = self.getLxy(muon)
0091           vz = muon.vz()
0092           theta_mu = 2*math.atan(math.exp(-muon.eta()))
0093           st1_z = (512.-lxy)/math.tan(theta_mu)+vz
0094           st1_r = 512.
0095           theta_st1 = math.atan2(st1_r,st1_z)
0096           eta_st1 = -math.log(math.tan(theta_st1/2.))
0097           return eta_st1
0098 
0099      def getSt2Phi(self,muon):  
0100           # calculate intersection of line and circle
0101           x1 = muon.vx()
0102           y1 = muon.vy()
0103           x2 = muon.vx() + muon.px()/(muon.px()**2+muon.py()**2)
0104           y2 = muon.vy() + muon.py()/(muon.px()**2+muon.py()**2)
0105           r = 512.
0106           dx = x2-x1
0107           dy = y2-y1
0108           dr = math.sqrt(dx**2+dy**2)
0109           D = x1*y2-x2*y1
0110           delta = (r**2)*(dr**2)-D**2
0111           if delta < 0:
0112                return math.atan2(y1, x1)
0113           # Two possible intersections = two possible phi values
0114           xP = (D*dy+math.copysign(1,dy)*dx*math.sqrt(delta))/dr**2
0115           xM = (D*dy-math.copysign(1,dy)*dx*math.sqrt(delta))/dr**2
0116           yP = (-D*dx+abs(dy)*math.sqrt(delta))/dr**2
0117           yM = (-D*dx-abs(dy)*math.sqrt(delta))/dr**2
0118           
0119           p1 = (xP, yP)
0120           p2 = (xM, yM)
0121           phi1 = math.atan2(yP,xP)
0122           phi2 = math.atan2(yM,xM)
0123           
0124           phi = min([phi1, phi2], key = lambda x: abs(self.deltaPhi(x, muon.phi()))) #probably a better way to select which intersection
0125           return phi   
0126 
0127 
0128      def process(self,gen,l1,dr=0.3,verbose=0):
0129           #first efficiency
0130           for g in gen:
0131                if verbose:
0132                     print("Gen Muon pt={pt} eta={eta} phi={phi} vxy={dxy}".format(pt=g.pt(),eta=g.eta(),phi=g.phi(),dxy=math.sqrt(g.vx()*g.vx()+g.vy()*g.vy())))
0133                self.histoData['genpt'].Fill(g.pt())
0134                if abs(g.eta())<0.83:
0135                     self.histoData['genptB'].Fill(g.pt())
0136                elif abs(g.eta())>0.83 and abs(g.eta())<1.2:      
0137                     self.histoData['genptO'].Fill(g.pt())
0138                else:
0139                     self.histoData['genptE'].Fill(g.pt())
0140                #Now let's process every threshold
0141                for t in self.thresholds:
0142                     if g.pt()>(t+10):
0143                          self.histoData['geneta'][t].Fill(g.eta())
0144                          self.histoData['genphi'][t].Fill(g.phi())
0145                          if self.getAcceptance(g):
0146                               self.histoData['genlxy'][t].Fill(self.getLxy(g))
0147                               
0148                     matched=[]
0149                     matchedDisplaced=[]
0150                     for mu in l1:                         
0151                          if mu.pt()<float(t):
0152                               continue
0153                          if(self.deltaR(g.eta(),g.phi(),mu.eta(),mu.phi()))<dr:
0154                             matched.append(mu)
0155                          if(self.deltaR(self.getSt2Eta(g),self.getSt2Phi(g),mu.eta(),mu.phi()))<dr:
0156                             matchedDisplaced.append(mu)
0157 
0158 #                    if len(matched)==0 and g.pt()>20 and self.prefix=='tk' and t==15:
0159 #                         import pdb;pdb.set_trace()
0160                     if len(matched)>0:
0161                          self.histoData['effpt'][t].Fill(g.pt())
0162                          if abs(g.eta())<0.83:
0163                               self.histoData['effptB'][t].Fill(g.pt())
0164                          elif abs(g.eta())>0.83 and abs(g.eta())<1.2:      
0165                               self.histoData['effptO'][t].Fill(g.pt())
0166                          else:
0167                               self.histoData['effptE'][t].Fill(g.pt())
0168                          if g.pt()>(t+10):
0169                               self.histoData['effeta'][t].Fill(g.eta())
0170                               self.histoData['effphi'][t].Fill(g.phi())
0171                          deltaPt=10000
0172                          best=None
0173                          for match in matched:
0174                               delta=abs(match.pt()-g.pt())
0175                               if delta<deltaPt:
0176                                    deltaPt=delta
0177                                    best=match
0178                          if deltaPt<10000:
0179                               self.histoData['resolution'].Fill(g.pt(),(best.pt()-g.pt())/g.pt())
0180 
0181                     if len(matchedDisplaced)>0:
0182                          if g.pt()>(t+10) and self.getAcceptance(g):
0183                               self.histoData['efflxy'][t].Fill(self.getLxy(g))
0184 
0185 
0186           #now rate
0187           maxElement=None
0188           maxPt=0
0189           for l in l1:
0190                if verbose:
0191                     print("{prefix} Muon pt={pt} eta={eta} phi={phi} stubs={stubs}".format(prefix=self.prefix,pt=l.pt(),eta=l.eta(),phi=l.phi(),stubs=len(l.stubs())))
0192                     for s in l.stubs():
0193                          print("-----> Associated Stub etaR={eta} phiR={phi} depthR={depth} coord1={coord1} coord2={coord2} q={q} ".format(eta=s.etaRegion(),phi=s.phiRegion(),depth=s.depthRegion(),coord1=s.offline_coord1(),coord2=s.offline_coord2(),q=s.quality()))
0194                if l.pt()>maxPt:
0195                     maxPt=l.pt()
0196                     maxElement=l
0197           if maxElement!=None:
0198                self.histoData['rate'].Fill(maxPt)
0199                if abs(maxElement.eta())<0.83:
0200                     self.histoData['rateBarrel'].Fill(maxPt)
0201                elif abs(maxElement.eta())>0.83 and abs(maxElement.eta())<1.2:
0202                     self.histoData['rateOverlap'].Fill(maxPt)
0203                else:
0204                     self.histoData['rateEndcap'].Fill(maxPt)
0205                for t in self.thresholds:
0206                     if maxPt>t:
0207                          self.histoData['rateeta'][t].Fill(maxElement.eta())                    
0208 
0209      def write(self,f):
0210           f.cd()
0211           self.histoData['genpt'].Write()
0212           self.histoData['genptB'].Write()
0213           self.histoData['genptO'].Write()
0214           self.histoData['genptE'].Write()
0215           self.histoData['resolution'].Write()
0216           c =self.histoData['rate'].GetCumulative(False)
0217           c.Scale(float(self.bunchfactor)/float(counter))
0218           c.Write(self.prefix+"_rate")     
0219           c =self.histoData['rateBarrel'].GetCumulative(False)
0220           c.Scale(float(self.bunchfactor)/float(counter))
0221           c.Write(self.prefix+"_rateBarrel")     
0222           c =self.histoData['rateOverlap'].GetCumulative(False)
0223           c.Scale(float(self.bunchfactor)/float(counter))
0224           c.Write(self.prefix+"_rateOverlap")     
0225           c =self.histoData['rateEndcap'].GetCumulative(False)
0226           c.Scale(float(self.bunchfactor)/float(counter))
0227           c.Write(self.prefix+"_rateEndcap")     
0228           for t in self.thresholds:
0229                c =self.histoData['rateeta'][t]
0230                c.Scale(float(self.bunchfactor)/float(counter))
0231                c.Write(self.prefix+"_rateeta_"+str(t))     
0232                g = ROOT.TGraphAsymmErrors(self.histoData['effpt'][t],self.histoData['genpt'])
0233                g.Write(self.prefix+"_eff_"+str(t))
0234                g = ROOT.TGraphAsymmErrors(self.histoData['effptB'][t],self.histoData['genptB'])
0235                g.Write(self.prefix+"_effB_"+str(t))
0236                g = ROOT.TGraphAsymmErrors(self.histoData['effptO'][t],self.histoData['genptO'])
0237                g.Write(self.prefix+"_effO_"+str(t))
0238                g = ROOT.TGraphAsymmErrors(self.histoData['effptE'][t],self.histoData['genptE'])
0239                g.Write(self.prefix+"_effE_"+str(t))
0240                g = ROOT.TGraphAsymmErrors(self.histoData['effeta'][t],self.histoData['geneta'][t])
0241                g.Write(self.prefix+"_effeta_"+str(t))
0242                g = ROOT.TGraphAsymmErrors(self.histoData['effphi'][t],self.histoData['genphi'][t])
0243                g.Write(self.prefix+"_effphi_"+str(t))
0244                g = ROOT.TGraphAsymmErrors(self.histoData['efflxy'][t],self.histoData['genlxy'][t])
0245                g.Write(self.prefix+"_efflxy_"+str(t))
0246                
0247           
0248 
0249 
0250 #HELPERS          
0251 def strAppend(s):
0252      return "root://cmsxrootd.fnal.gov/"+s
0253 
0254 def fetchSTA(event,tag,etamax=3.0):
0255      phiSeg2    = Handle  ('vector<l1t::SAMuon>')
0256      event.getByLabel(tag,phiSeg2)
0257      return list(filter(lambda x: abs(x.eta())<etamax,phiSeg2.product()))
0258 def fetchTPS(event,tag,etamax=3.0):
0259      phiSeg2    = Handle  ('vector<l1t::TrackerMuon>')
0260      event.getByLabel(tag,phiSeg2)
0261      return list(filter(lambda x: abs(x.eta())<etamax,phiSeg2.product()))
0262 def fetchStubs(event,tag):
0263      phiSeg2    = Handle  ('vector<l1t::MuonStub>')
0264      event.getByLabel(tag,phiSeg2)
0265      return phiSeg2.product()
0266 
0267 def fetchGEN(event,etaMax=3.0):
0268      genH  = Handle  ('vector<reco::GenParticle>')
0269      event.getByLabel('genParticles',genH)
0270      genMuons=list(filter(lambda x: abs(x.pdgId())==13 and x.status()==1 and abs(x.eta())<etaMax,genH.product()))
0271      return genMuons
0272 
0273 def EOSls(path):
0274     print('eos root://cmseos.fnal.gov/ find -name "*root" %s' % path )
0275     p = subprocess.Popen('eos root://cmseos.fnal.gov/ find -name "*root" %s' % path, 
0276                          stdout = subprocess.PIPE, stderr = subprocess.PIPE,
0277                          shell=True)
0278     stdout, stderr = p.communicate()
0279     out = ["root://cmseos.fnal.gov/"+i.decode('UTF-8') for i in stdout.split()]
0280     return out
0281 #DATASET
0282 # from samples200 import *
0283 # toProcess = dy200
0284 # files=[]
0285 # for p in toProcess:
0286      # files.append(strAppend(p))
0287 # events=Events(files)
0288 # tag='DY_v1'
0289 # tag='disp200'
0290 
0291 parser = argparse.ArgumentParser(description='Process some integers.')
0292 parser.add_argument('--tag', default="Tau", help="fdf")
0293 parser.add_argument('--prod', default="gmtMuons", help="fdf")
0294 args = parser.parse_args()
0295 tag = args.tag
0296 samples = {
0297     "MB_org" : "/eos/uscms//store/group/lpctrig/benwu/GMT_Ntupler/Spring22_GMToriginal_v2/MinBias_TuneCP5_14TeV-pythia8/PHASEII_MinBias/",
0298     "DY_org" : "/eos/uscms//store/group/lpctrig/benwu/GMT_Ntupler/Spring22_GMToriginal_v2/DYToLL_M-50_TuneCP5_14TeV-pythia8/PHASEII_DYToLL/",
0299     "MB_v1" : "/eos/uscms//store/group/lpctrig/benwu/GMT_Ntupler/Spring22_GMT_v3/MinBias_TuneCP5_14TeV-pythia8/PHASEII_MinBias/",
0300     "DY_v1" : "/eos/uscms//store/group/lpctrig/benwu/GMT_Ntupler/Spring22_GMT_v3/DYToLL_M-50_TuneCP5_14TeV-pythia8/PHASEII_DYToLL/",
0301 }
0302 flist = EOSls(samples[tag])
0303 events= Events(flist)
0304 # events=Events(['file:/uscmst1b_scratch/lpc1/3DayLifetime/bachtis/gmt.root'])
0305 #events=Events(['file:gmt.root'])
0306 # events=Events(['file:reprocess.root'])
0307 
0308 
0309 #ANALYSIS SETUP
0310 verbose=0
0311 saAnalyzer = muon_trigger_analyzer('sta_prompt')
0312 kmtfAnalyzer_p = muon_trigger_analyzer('kmtf_prompt')
0313 kmtfAnalyzer_d = muon_trigger_analyzer('kmtf_disp')
0314 tkAnalyzer_1st = muon_trigger_analyzer('tk_1st')
0315 tkAnalyzer_2st = muon_trigger_analyzer('tk_2st')
0316 tkAnalyzer_2m = muon_trigger_analyzer('tk_2m')
0317 tkAnalyzer_3m = muon_trigger_analyzer('tk_3m')
0318 tkAnalyzer_4m = muon_trigger_analyzer('tk_4m')
0319 
0320 #EVENT LOOP
0321 
0322 counter=0;
0323 for event in events:
0324      if verbose:
0325           print('EVENT {}'.format(counter))
0326      gen=fetchGEN(event,2.4)
0327      genKMTF=fetchGEN(event,0.83)    
0328      sa=fetchSTA(event,'gmtSAMuons:prompt',2.5)
0329      kmtf_p=fetchSTA(event,'gmtKMTFMuons:prompt',2.5)
0330      kmtf_d=fetchSTA(event,'gmtKMTFMuons:displaced',2.5)
0331      tps=fetchTPS(event,'gmtTkMuons',2.5)
0332      # sa=fetchSTA(event,'l1tSAMuonsGmt:prompt',2.5)
0333      # kmtf_p=fetchSTA(event,'l1tKMTFMuonsGmt:prompt',2.5)
0334      # kmtf_d=fetchSTA(event,'l1tKMTFMuonsGmt:displaced',2.5)
0335      # tps=fetchTPS(event,'l1tTkMuonsGmt',2.5)
0336      tps_2st=list(filter(lambda x: x.stubs().size()>1,tps))
0337      tps_2m=list(filter(lambda x: x.numberOfMatches()>1,tps))
0338      tps_3m=list(filter(lambda x: x.numberOfMatches()>2,tps))
0339      tps_4m=list(filter(lambda x: x.numberOfMatches()>3,tps))
0340      #analyze
0341      saAnalyzer.process(gen,sa,0.6,verbose)
0342      kmtfAnalyzer_p.process(genKMTF,kmtf_p,0.6,verbose)
0343      kmtfAnalyzer_d.process(genKMTF,kmtf_d,0.6,verbose)
0344      tkAnalyzer_1st.process(gen,tps,0.2,verbose)
0345      tkAnalyzer_2st.process(gen,tps_2st,0.2,0)
0346      tkAnalyzer_2m.process(gen,tps_2m,0.2,0)
0347      tkAnalyzer_3m.process(gen,tps_3m,0.2,0)
0348      tkAnalyzer_4m.process(gen,tps_4m,0.2,0)
0349 
0350 
0351 
0352 #     if verbose==1 and counter==69:
0353 #          import pdb;pdb.set_trace()
0354      #counter 
0355 #     if counter==100000:
0356 #          break
0357 
0358      if (counter %10000 ==0):
0359           print(counter)
0360      counter=counter+1
0361 
0362 f=ROOT.TFile("gmtAnalysis_{tag}.root".format(tag=tag),"RECREATE")
0363 f.cd()
0364 saAnalyzer.write(f)
0365 kmtfAnalyzer_p.write(f)
0366 kmtfAnalyzer_d.write(f)
0367 tkAnalyzer_1st.write(f)
0368 tkAnalyzer_2st.write(f)
0369 tkAnalyzer_2m.write(f)
0370 tkAnalyzer_3m.write(f)
0371 tkAnalyzer_4m.write(f)
0372 f.Close()