Utilities/scripts/pileupDistInMC.py

0001 #! /usr/bin/env python3
0002
0003 from __future__ import print_function
0004 import optparse
0005 import re
0006 from pprint import pprint
0007
0008 commentRE = re.compile (r'#.*$')
0009
0010 if __name__ == "__main__":
0011     parser = optparse.OptionParser ("Usage: %prog file1.root [file2.root...]")
0012     parser.add_option ('--loadFromFile', dest='loadFromFile', default=[],
0013                        type='string',
0014                        action='append',
0015                        help="Name of text file containing filenames" )
0016     parser.add_option ('--prefix', dest='prefix', type='string',
0017                        default='',
0018                        help="Prefix to add to files" )
0019
0020     parser.add_option ('--bx', dest='bx', type='int',
0021                        default='0',
0022                        help="Bunch crossing to check (0 = in-time)" )
0023     (options, args) = parser.parse_args()
0024     import ROOT # stupid ROOT takes the arugments error
0025     from DataFormats.FWLite import Events, Handle
0026
0027     listOfFiles = args[:]
0028     for filename in options.loadFromFile:
0029         source = open (filename, 'r')
0030         for line in source:
0031             line = commentRE.sub ('', line).strip() # get rid of comments
0032             if not line:
0033                 # don't bother with blank lines
0034                 continue
0035             listOfFiles.append (line)
0036         source.close()
0037     if options.prefix:
0038         oldList = listOfFiles
0039         listOfFiles = []
0040         for name in oldList:
0041             listOfFiles.append( options.prefix + name )
0042
0043     if not listOfFiles:
0044         raise RuntimeError("You have not provided any files")
0045
0046     events = Events (listOfFiles)
0047
0048     handle = Handle('vector<PileupSummaryInfo>')
0049     label  = ('addPileupInfo')
0050
0051     ROOT.gROOT.SetBatch()        # don't pop up canvases
0052
0053     # loop over events
0054     countDict = {}
0055     total = 0.
0056     for event in events:
0057         event.getByLabel (label, handle)
0058         pileups = handle.product()
0059         for pileup in pileups:
0060             if pileup.getBunchCrossing() == options.bx:
0061                 break
0062             if pileup == pileups[-1] and len(pileups)>1 :
0063                 raise RuntimeError("Requested BX not found in file")
0064
0065         num = pileup.getPU_NumInteractions()
0066         total += 1
0067         if num not in countDict:
0068             countDict[num] = 1
0069         else:
0070             countDict[num] += 1
0071
0072     print("total", int(total), "\ncounts:")
0073     pprint (countDict, width=1)
0074     print("normalized:")
0075
0076     renormDict = {}
0077     for key, count in countDict.items():
0078         renormDict[key] = count / total
0079     pprint (renormDict)
0080