Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2023-03-17 10:58:57

0001 #!/usr/bin/env python3
0002 
0003 """Usage: test_fastHaddMerge.py [-t produce|check] [-n files]
0004 
0005 Produce a custom number of identical ROOT files and check that their
0006 final merged output matches what is expected.
0007 """
0008 from __future__ import print_function
0009 
0010 from builtins import range
0011 from optparse import OptionParser
0012 import sys
0013 import subprocess
0014 import re
0015 
0016 word2num = {'One': 1, 'Two': 2, 'Ten': 10, 'Twenty': 20, 'Fifty': 50}
0017 
0018 class Histo(object):
0019     def __init__(self, name, entries,
0020                  bins, xmin, xmax,
0021                  value, folder):
0022         self.name_ = name
0023         self.bins_ = bins
0024         self.xmin_ = xmin
0025         self.xmax_ = xmax
0026         self.entries_ = entries
0027         self.value_ = value
0028         self.folder_ = folder
0029 
0030     def book_and_fill(self, tfile):
0031         import ROOT as r
0032         r.gDirectory.cd('/')
0033         fullpath = ''
0034         for level in self.folder_.split('/'):
0035             fullpath += '/%s' % level
0036             if not level == '':
0037                 if not r.gDirectory.GetDirectory(level):
0038                     r.gDirectory.mkdir(level)
0039                 r.gDirectory.cd(fullpath)
0040         histo = r.TH1F(self.name_, self.name_,
0041                        self.bins_, self.xmin_, self.xmax_)
0042         for e in range(0, self.entries_):
0043             histo.Fill(self.value_)
0044         histo.Write()
0045 
0046 class FileProducer(object):
0047     def __init__(self, prefix, numFiles, folders, histo_per_folder):
0048         self.prefix_ = prefix
0049         self.numFiles_ = numFiles
0050         self.folders_ = folders
0051         self.histo_per_folder_ = histo_per_folder
0052         self.histos_ = []
0053         assert numFiles > 0
0054 
0055     def checkCumulative(self, rootfile):
0056         import ROOT as r
0057         self.prepareHistos()
0058         f = r.TFile(rootfile)
0059         num = 0
0060         for histo in self.histos_:
0061             num += 1
0062             if num%10000 == 0:
0063                 sys.stdout.write('.')
0064                 sys.stdout.flush()
0065             h = f.Get(histo.folder_+'/'+histo.name_)
0066             h.SetDirectory(0)
0067             assert h
0068             m = re.match('(.*)Entr.*_(\d+)$', h.GetName())
0069             assert m.groups()
0070             assert h.GetEntries() == word2num[m.group(1)] * self.numFiles_
0071             assert h.GetMean() == float(m.group(2))
0072         print()
0073         f.Close()
0074         print()
0075 
0076     def createIdenticalFiles(self):
0077         import ROOT as r
0078         self.prepareHistos()
0079         f = r.TFile("%s_%d.root" % (self.prefix_, 0), "RECREATE")
0080         num = 0
0081         for h in self.histos_:
0082             num = num + 1
0083             if num%1000 == 0:
0084                 sys.stdout.write(".")
0085                 sys.stdout.flush()
0086             h.book_and_fill(f)
0087         f.Write()
0088         f.Close()
0089         print('Wrote %d histograms in %d folders' % (len(self.histos_), len(self.folders_)))
0090         for i in range(1, self.numFiles_):
0091             subprocess.getoutput("cp %s_0.root %s_%d.root" % (self.prefix_,
0092                                                             self.prefix_,
0093                                                             i))
0094             print('x')
0095 
0096     def prepareHistos(self):
0097         for folder in self.folders_:
0098             sys.stdout.write("|")
0099             sys.stdout.flush()
0100             for i in range(0, self.histo_per_folder_):
0101                 if i%100 == 0:
0102                     sys.stdout.write(".")
0103                     sys.stdout.flush()
0104                 self.histos_.append(Histo("OneEntry_%d" %i,
0105                                           1, i+1, 0, i+1, i, folder))
0106                 self.histos_.append(Histo("TwoEntries_%d" %i,
0107                                           2, i+1, 0, i+1, i, folder))
0108                 self.histos_.append(Histo("TenEntries_%d" %i,
0109                                           10, i+1, 0, i+1, i, folder))
0110                 self.histos_.append(Histo("TwentyEntries_%d" %i,
0111                                           20, i+1, 0, i+1, i, folder))
0112                 self.histos_.append(Histo("FiftyEntries_%d" %i,
0113                                           50, i+1, 0, i+1, i, folder))
0114         print()
0115 
0116 op = OptionParser(usage = __doc__)
0117 op.add_option("-a", "--action", dest = "action",
0118               type = "string", action = "store", metavar = "ACTION",
0119               default = "produce",
0120               help = "Either produce or check ROOT file(s).")
0121 op.add_option("-c", "--check", dest = "file_to_check",
0122               type = "string", action = "store", metavar = "FILE",
0123               default = '',
0124               help = "Check the content of FILE against expectations.")
0125 op.add_option("-n", "--numfiles", dest = "numfiles",
0126               type = "int", action = "store", metavar = "NUM",
0127               default = 10, help = "Create NUM identical files")
0128 options, args = op.parse_args()
0129 
0130 if __name__ == '__main__':
0131     fp = FileProducer("MergePBTest", options.numfiles, ["/First_level",
0132                                                         "/Second/Level",
0133                                                         "/Third/Level/Folder",
0134                                                         "/Fourth/Of/Many/Folders",
0135                                                         "/Pixel/", "/Pixel/A", "/Pixel/A/B", "/Pixel/A/B/C", "/Pixel/A/B/C/D", "/Pixel/A/B/C/D/E",
0136                                                         "/SiStrip/", "/SiStrip/A", "/SiStrip/A/B", "/SiStrip/A/B/C", "/SiStrip/A/B/C/D", "/SiStrip/A/B/C/D/E",
0137                                                         "/RPC/", "/RPC/A", "/RPC/A/B", "/RPC/A/B/C", "/RPC/A/B/C/D", "/RPC/A/B/C/D/E",
0138                                                         "/HLT/", "/HLT/A", "/HLT/A/B", "/HLT/A/B/C", "/HLT/A/B/C/D", "/HLT/A/B/C/D/E",
0139                                                         "/EcalBarrel/", "/EcalBarrel/A", "/EcalBarrel/A/B", "/EcalBarrel/A/B/C", "/EcalBarrel/A/B/C/D", "/EcalBarrel/A/B/C/D/E",
0140                                                         "/EcalEndcap/", "/EcalEndcap/A", "/EcalEndcap/A/B", "/EcalEndcap/A/B/C", "/EcalEndcap/A/B/C/D", "/EcalEndcap/A/B/C/D/E",
0141                                                         "/Tracking/", "/Tracking/A", "/Tracking/A/B", "/Tracking/A/B/C", "/Tracking/A/B/C/D", "/Tracking/A/B/C/D/E",
0142                                                         "/Muon/", "/Muon/A", "/Muon/A/B", "/Muon/A/B/C", "/Muon/A/B/C/D", "/Muon/A/B/C/D/E",
0143                                                         "/EGamma/", "/EGamma/A", "/EGamma/A/B", "/EGamma/A/B/C", "/EGamma/A/B/C/D", "/EGamma/A/B/C/D/E",
0144                                                         "/Tau/", "/Tau/A", "/Tau/A/B", "/Tau/A/B/C", "/Tau/A/B/C/D", "/Tau/A/B/C/D/E"],
0145                       100)
0146     if options.action == 'produce':
0147         fp.createIdenticalFiles()
0148     else:
0149         if not options.action == 'check':
0150             print("Option -a|--action takes only 'produce|check' options.", file=sys.stderr)
0151             sys.exit(1)
0152         else:
0153             if options.file_to_check == '':
0154                 print("Option -c|--check required to check the content of a file.", file=sys.stderr)
0155                 sys.exit(1)
0156             fp.checkCumulative(options.file_to_check)
0157 
0158 # Local Variables:
0159 # show-trailing-whitespace: t
0160 # truncate-lines: t
0161 # End: