Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2024-11-26 02:34:12

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