postprocessing/examples/exampleModule.py

0001 # This is an example of a NanoAODTools Module to add one variable to nanoAODs.
0002 # Note that:
0003 # -the new variable will be available for use in the subsequent modules
0004 # -it is possible to update the value for existing variables
0005 #
0006 # Example of using from command line:
0007 # nano_postproc.py outDir /eos/cms/store/user/andrey/f.root -I PhysicsTools.NanoAODTools.postprocessing.examples.exampleModule exampleModuleConstr
0008 #
0009 # Example of running in a python script: see test/example_postproc.py
0010 #
0011
0012 from PhysicsTools.NanoAODTools.postprocessing.framework.datamodel import Collection
0013 from PhysicsTools.NanoAODTools.postprocessing.framework.eventloop import Module
0014 import ROOT
0015 ROOT.PyConfig.IgnoreCommandLineOptions = True
0016
0017
0018 class exampleProducer(Module):
0019     def __init__(self, jetSelection):
0020         self.jetSel = jetSelection
0021         pass
0022
0023     def beginJob(self):
0024         pass
0025
0026     def endJob(self):
0027         pass
0028
0029     def beginFile(self, inputFile, outputFile, inputTree, wrappedOutputTree):
0030         self.out = wrappedOutputTree
0031         self.out.branch("EventMass", "F")
0032
0033     def endFile(self, inputFile, outputFile, inputTree, wrappedOutputTree):
0034         pass
0035
0036     def analyze(self, event):
0037         """process event, return True (go to next module) or False (fail, go to next event)"""
0038         electrons = Collection(event, "Electron")
0039         muons = Collection(event, "Muon")
0040         jets = Collection(event, "Jet")
0041         eventSum = ROOT.TLorentzVector()
0042         for lep in muons:
0043             eventSum += lep.p4()
0044         for lep in electrons:
0045             eventSum += lep.p4()
0046         for j in filter(self.jetSel, jets):
0047             eventSum += j.p4()
0048         self.out.fillBranch("EventMass", eventSum.M())
0049         return True
0050
0051
0052 # define modules using the syntax 'name = lambda : constructor' to avoid having them loaded when not needed
0053
0054 exampleModuleConstr = lambda: exampleProducer(jetSelection=lambda j: j.pt > 30)