postprocessing/examples/mhtjuProducerCpp.py

0001 # Example of calling a C++ helper from a Module.
0002 #
0003 # Run with:
0004 # nano_postproc.py outDir /eos/cms/store/user/andrey/f.root -I PhysicsTools.NanoAODTools.postprocessing.examples.mhtjuProducerCpp mhtju
0005
0006 from PhysicsTools.NanoAODTools.postprocessing.framework.eventloop import Module
0007 import ROOT
0008 import os
0009 ROOT.PyConfig.IgnoreCommandLineOptions = True
0010
0011
0012 # MHT producer, unclean jets only (no lepton overlap cleaning, no jet selection)
0013 class mhtjuProducerCpp(Module):
0014     def __init__(self):
0015         base = os.getenv("NANOAODTOOLS_BASE")
0016         if base:
0017             # Running in standalone mode: compile the C++ helper
0018             if "/MhtjuProducerCppWorker_cc.so" not in ROOT.gSystem.GetLibraries():
0019                 print("Load C++ MhtjuProducerCppWorker worker module")
0020                 ROOT.gROOT.ProcessLine(
0021                     ".L %s/test/examples/MhtjuProducerCppWorker.cc+O" % base)
0022         else:
0023             # Load the helper from the CMSSW compiled. This is not required if
0024             # dictionaries for the helper are generated with classes_def.xml and
0025             # classes.h
0026             base = "%s/src/PhysicsTools/NanoAODTools" % os.getenv("CMSSW_BASE")
0027             ROOT.gSystem.Load("libPhysicsToolsNanoAODToolsTest.so")
0028             ROOT.gROOT.ProcessLine(".L %s/test/examples/MhtjuProducerCppWorker.h" % base)
0029         self.worker = ROOT.MhtjuProducerCppWorker()
0030         pass
0031
0032     def beginJob(self):
0033         pass
0034
0035     def endJob(self):
0036         pass
0037
0038     def beginFile(self, inputFile, outputFile, inputTree, wrappedOutputTree):
0039         self.initReaders(inputTree)  # initReaders must be called in beginFile
0040         self.out = wrappedOutputTree
0041         self.out.branch("MHTju_pt", "F")
0042         self.out.branch("MHTju_phi", "F")
0043
0044     def endFile(self, inputFile, outputFile, inputTree, wrappedOutputTree):
0045         pass
0046
0047     # this function gets the pointers to Value and ArrayReaders and sets
0048     # them in the C++ worker class
0049     def initReaders(self, tree):
0050         self.nJet = tree.valueReader("nJet")
0051         self.Jet_pt = tree.arrayReader("Jet_pt")
0052         self.Jet_phi = tree.arrayReader("Jet_phi")
0053         self.worker.setJets(self.nJet, self.Jet_pt, self.Jet_phi)
0054         # self._ttreereaderversion must be set AFTER all calls to
0055         # tree.valueReader or tree.arrayReader
0056         self._ttreereaderversion = tree._ttreereaderversion
0057
0058     def analyze(self, event):
0059         """process event, return True (go to next module) or False (fail,
0060         go to next event)"""
0061
0062         # do this check at every event, as other modules might have read
0063         # further branches
0064         if event._tree._ttreereaderversion > self._ttreereaderversion:
0065             self.initReaders(event._tree)
0066         # do NOT access other branches in python between the check/call to
0067         # initReaders and the call to C++ worker code
0068         output = self.worker.getHT()
0069
0070         self.out.fillBranch("MHTju_pt", output[0])
0071         self.out.fillBranch("MHTju_phi", -output[1])  # note the minus
0072         return True
0073
0074
0075 # define modules using the syntax 'name = lambda : constructor' to avoid
0076 # having them loaded when not needed
0077
0078 mhtju = lambda: mhtjuProducerCpp()