analyzers/examples/ntuple.py

0001 #!/bin/env python
0002
0003 def var( tree, varName, type=float ):
0004     tree.var(varName, type)
0005
0006 def fill( tree, varName, value ):
0007     tree.fill( varName, value )
0008
0009 # event information
0010
0011 def bookEvent(tree):
0012     var(tree, 'run')
0013     var(tree, 'lumi')
0014     var(tree, 'event')
0015
0016 def fillEvent(tree, event):
0017     fill(tree, 'run', event.run)
0018     fill(tree, 'lumi', event.lumi)
0019     fill(tree, 'event', event.eventId)
0020
0021
0022 # simple particle
0023
0024 def bookParticle( tree, pName ):
0025     var(tree, '{pName}_pdgid'.format(pName=pName))
0026     var(tree, '{pName}_e'.format(pName=pName))
0027     var(tree, '{pName}_pt'.format(pName=pName))
0028     var(tree, '{pName}_eta'.format(pName=pName))
0029     var(tree, '{pName}_phi'.format(pName=pName))
0030     var(tree, '{pName}_m'.format(pName=pName))
0031     var(tree, '{pName}_q'.format(pName=pName))
0032
0033 def fillParticle( tree, pName, particle ):
0034     fill(tree, '{pName}_pdgid'.format(pName=pName), particle.pdgId() )
0035     fill(tree, '{pName}_e'.format(pName=pName), particle.energy() )
0036     fill(tree, '{pName}_pt'.format(pName=pName), particle.pt() )
0037     fill(tree, '{pName}_eta'.format(pName=pName), particle.eta() )
0038     fill(tree, '{pName}_phi'.format(pName=pName), particle.phi() )
0039     fill(tree, '{pName}_m'.format(pName=pName), particle.mass() )
0040     fill(tree, '{pName}_q'.format(pName=pName), particle.charge() )
0041
0042 def bookMet(tree, pName):
0043     var(tree, '{pName}_pt'.format(pName=pName))
0044     var(tree, '{pName}_phi'.format(pName=pName))
0045     var(tree, '{pName}_sumet'.format(pName=pName))
0046
0047 def fillMet(tree, pName, met):
0048     fill(tree, '{pName}_pt'.format(pName=pName), met.pt())
0049     fill(tree, '{pName}_phi'.format(pName=pName), met.phi())
0050     fill(tree, '{pName}_sumet'.format(pName=pName), met.sumEt())
0051
0052 def bookGenTau(tree, pName, pfdiscs, calodiscs):
0053     bookJet(tree, pName)
0054     bookTau(tree, '{pName}_calo'.format(pName=pName), calodiscs)
0055     bookTau(tree, '{pName}_pf'.format(pName=pName), pfdiscs)
0056     bookJet(tree, '{pName}_pfjet'.format(pName=pName))
0057
0058 def fillGenTau(tree, pName, tau):
0059     fillJet(tree, pName, tau)
0060     fillTau(tree, '{pName}_calo'.format(pName=pName), tau.match_calo)
0061     fillTau(tree, '{pName}_pf'.format(pName=pName), tau.match_pf)
0062     fillJet(tree, '{pName}_pfjet'.format(pName=pName), tau.match_pfjet)
0063
0064
0065 def bookTau(tree, pName, discNames):
0066     bookParticle(tree, pName)
0067     var(tree, '{pName}_nsigcharged'.format(pName=pName))
0068     var(tree, '{pName}_isolation'.format(pName=pName))
0069     for discName in discNames:
0070         var(tree, '{pName}_{disc}'.format(pName=pName,
0071                                           disc=discName))
0072
0073 def fillTau(tree, pName, tau):
0074     if not tau: return
0075     fillParticle(tree, pName, tau)
0076     fill(tree, '{pName}_nsigcharged'.format(pName=pName), len(tau.signalCharged()))
0077     fill(tree, '{pName}_isolation'.format(pName=pName), tau.isolation())
0078     for discName, value in tau.discs.items():
0079         fill(tree, '{pName}_{disc}'.format(pName=pName,
0080                                            disc=discName), value)
0081
0082
0083 # jet
0084
0085 def bookComponent( tree, pName ):
0086     var(tree, '{pName}_e'.format(pName=pName))
0087     var(tree, '{pName}_pt'.format(pName=pName))
0088     var(tree, '{pName}_num'.format(pName=pName))
0089
0090 def fillComponent(tree, pName, component):
0091     fill(tree, '{pName}_e'.format(pName=pName), component.e() )
0092     fill(tree, '{pName}_pt'.format(pName=pName), component.pt() )
0093     fill(tree, '{pName}_num'.format(pName=pName), component.num() )
0094
0095
0096 pdgids = [211, 22, 130, 11, 13]
0097
0098 def bookJet( tree, pName ):
0099     bookParticle(tree, pName )
0100     for pdgid in pdgids:
0101         bookComponent(tree, '{pName}_{pdgid:d}'.format(pName=pName, pdgid=pdgid))
0102     # var(tree, '{pName}_npart'.format(pName=pName))
0103
0104 def fillJet( tree, pName, jet ):
0105     if not jet: return
0106     fillParticle(tree, pName, jet )
0107     for pdgid in pdgids:
0108         component = jet.constituents.get(pdgid, None)
0109         if component is not None:
0110             fillComponent(tree,
0111                           '{pName}_{pdgid:d}'.format(pName=pName, pdgid=pdgid),
0112                           component )
0113         else:
0114             import pdb; pdb.set_trace()
0115             print(jet)
0116