File indexing completed on 2024-11-26 02:34:12
0001
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
0158
0159
0160