Utilities/scripts/pileupDistInMC.py

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