analyzers/examples/ntuple.py

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