File indexing completed on 2024-09-27 07:02:50
0001
0002
0003 from array import array
0004 import os
0005
0006 import ROOT
0007 try:
0008 ROOT.PyConfig.IgnoreCommandLineOptions = True
0009 except AttributeError:
0010 pass
0011 ROOT.gSystem.Load('libeicsmear')
0012 ROOT.gROOT.SetBatch(True)
0013
0014 def parse():
0015 """Parse command line arguments and return in an argparse.Namespace."""
0016 from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
0017 parser = ArgumentParser(formatter_class=ArgumentDefaultsHelpFormatter)
0018 parser.add_argument('file', nargs='+', help='input ASCII file')
0019 parser.add_argument('-o', '--outdir', default='.',
0020 help='output directory')
0021 parser.add_argument('-f', '--format', default='pdf', choices=['ps', 'pdf'],
0022 help='output file format')
0023 parser.add_argument('-n', '--nevents', default=-1, type=int,
0024 help='number of events per file')
0025 return parser.parse_args()
0026
0027 def build(filename, outdir='.', nevents=-1):
0028 """Build a ROOT tree from the named text file.
0029
0030 Returns the name of the resultant ROOT file.
0031
0032 """
0033 base, extension = os.path.splitext(filename)
0034 ROOT.BuildTree(filename, outdir, nevents)
0035 return '.'.join([base, 'root'])
0036
0037 def histogram(typename, *args, **kwargs):
0038 """Return a histogram of the requested type.
0039
0040 args -- normal histogram constructor arguments
0041 kwargs -- additional attributes to attach to the histogram
0042
0043 e.g. to create a TH1D with a 'cut' attribute:
0044
0045 histogram(ROOT.TH1D, 'myhist', '', 10, 0, 10, cut='x>10')
0046
0047 """
0048 h = typename(*args)
0049 attributes = {'select': h.GetName(), 'cut': '', 'log': ''}
0050 attributes.update(kwargs)
0051 for name, value in attributes.iteritems():
0052 setattr(h, name, value)
0053 h.SetOption(attributes.get('opt', ''))
0054 return h
0055
0056 def th1d(*args, **kwargs):
0057 """Return a TH1D"""
0058 return histogram(ROOT.TH1D, *args, **kwargs)
0059
0060 def th2d(*args, **kwargs):
0061 """Return a TH2D"""
0062 return histogram(ROOT.TH2D, *args, **kwargs)
0063
0064 def fill_histogram_from_tree(self, tree):
0065 """Fill a histogram from a tree.
0066
0067 Uses the histogram's 'select' and 'cut' attributes to determine
0068 what is plotted.
0069
0070 """
0071 try:
0072 counts = tree.Project(self.GetName(), self.select, self.cut)
0073 except:
0074 print 'Error filling histogram', self.GetName(),\
0075 'from tree, histogram will not be filled'
0076
0077
0078 ROOT.TH1.FillTree = fill_histogram_from_tree
0079
0080 class Histograms(object):
0081 """A collection of basic histograms."""
0082
0083 def name(self, string):
0084 """Return the string prepended by this object's file basename."""
0085 return '_'.join([self.base, string])
0086
0087 def __init__(self, rootfile):
0088 """Initialise histograms and associate them with a ROOT file."""
0089 rootfile.cd()
0090 self.file = rootfile
0091 self.base, extension = os.path.splitext(file.GetName())
0092 self.out = '.'.join([self.base, 'pdf'])
0093 name = self.name
0094 self.window = ROOT.TCanvas(name('canvas'), '', 1, 1, 800, 800)
0095 self.page(open=True)
0096
0097 self.histograms = [
0098 th1d(name('beamlepton'), 'Beam lepton ID', 300, -3000., 3000.,
0099 select='BeamLepton().Id().Code()'),
0100 th1d(name('beamhadron'), 'Beam hadron ID', 300, -3000., 3000.,
0101 select='BeamHadron().Id().Code()'),
0102 th1d(name('scattered'), 'Scattered lepton ID', 300, -3000., 3000.,
0103 select='ScatteredLepton().Id().Code()'),
0104 th1d(name('boson'), 'Exchange boson ID', 300, -3000., 3000.,
0105 select='ExchangeBoson().Id().Code()'),
0106 th1d(name('nu'), 'nu', 25, 1., 5., select='log10(nu)'),
0107 th2d(name('Q2x'), '$Q^{2}$ vs. $x_{Bj}$', 25, -5., 0., 25, -1., 4.,
0108 select='log10(QSquared):log10(x)', opt='colz'),
0109 th1d(name('xf'), '$x_{F}$', 55, -1.1, 1.1, select='xFeynman',
0110 cut='KS==1 && id!=ScatteredLepton().Id().Code()'),
0111 th1d(name('z'), '$z$', 60, -0.1, 1.1, select='z',
0112 cut='KS==1 && id!=ScatteredLepton().Id().Code()', log='y'),
0113 th1d(name('pt'), '$p_{T}$ vs. virtual photon', 60, 0., 3.,
0114 select='ptVsGamma', log='y'),
0115 th1d(name('ptot'), 'Total final-state momentum', 100, 0., 300.,
0116 select='FinalStateMomentum().P()'),
0117 th1d(name('ctot'), 'Total final-state charge', 100, -2., 2.,
0118 select='FinalStateCharge()'),
0119 th2d(name('pthb'), 'Get4VectorInHadronBosonFrame().Pt():ptVsGamma',
0120 50, 0., 10., 50, 0., 10.,
0121 select='Get4VectorInHadronBosonFrame().Pt():ptVsGamma',
0122 opt='colz'),
0123 th2d(name('mass'), 'Photon mass$^{2}$ vs. $Q^{2}$',
0124 25, -1., 4., 25, -1., 4.,
0125 select='log10(-(ExchangeBoson().Get4Vector().M2())):'
0126 'log10(QSquared)',
0127 opt='colz')
0128 ]
0129
0130 for method in ['', 'JB', 'DA']:
0131 def var(string):
0132 """Insert method into a string."""
0133 if '{' in string:
0134 return string.format(method)
0135 else:
0136 return ''.join([string, method])
0137 self.histograms += [
0138 th1d(name(var('x')), var('$x_{{Bj}}{}$'), 25, -5., 0.,
0139 select=var('log10(x{})')),
0140 th1d(name(var('Q2')), var('$Q^{{2}}{}$'), 25, -1., 4.,
0141 select=var('log10(QSquared)'), log='y'),
0142 th1d(name(var('y')), var('$y{}$'), 25, -0.1, 1.1,
0143 select=var('y')),
0144 th1d(name(var('W2')), var('$W^{{2}}{}$'), 25, 1., 6.,
0145 select=var('log10(WSquared)'))
0146 ]
0147
0148 self.histograms += [
0149 th1d(name('number'), 'Event number', 100, 0.,
0150 file.EICTree.GetEntries(), select='number'),
0151 th1d(name('process'), 'Process', 200, 0., 200., select='process'),
0152 th1d(name('nTracks'), 'Number of tracks', 200, 0., 200.,
0153 select='nTracks'),
0154 th2d(name('tracks'), 'particles.GetEntries() vs. nTracks',
0155 200, 0., 200., 200, 0., 200.,
0156 select='@particles.GetEntries():nTracks', opt='colz'),
0157 th1d(name('Ein'), 'ELeptonInNucl', 25, 0., 12000.,
0158 select='ELeptonInNucl'),
0159 th1d(name('Eout'), 'ELeptonOutNucl', 25, 0., 12000.,
0160 select='ELeptonOutNucl'),
0161 th2d(name('hadronic'),
0162 'HadronicFinalStateMomentum().M2():WSquared',
0163 25, 1., 6., 25, 1., 6.,
0164 select='log10(HadronicFinalStateMomentum().M2()):'
0165 'log10(WSquared)', opt='colz')
0166 ]
0167
0168 def page(self, open=False, close=False, log=''):
0169 """Print the current contents of the window."""
0170 if open:
0171 self.window.Print(self.out + '[')
0172 elif close:
0173 self.window.Print(self.out + ']')
0174 else:
0175 self.window.SetLogx('x' in log)
0176 self.window.SetLogy('y' in log)
0177 self.window.SetLogz('z' in log)
0178 self.window.Print(self.out)
0179
0180 def make(self):
0181 """Fill and print all histograms."""
0182 self.file.cd()
0183 for histogram in self.histograms:
0184 histogram.FillTree(self.file.EICTree)
0185 self.window.cd()
0186 histogram.Draw()
0187 self.page(log=histogram.log)
0188
0189 def close(self):
0190 """Cloes the ROOT file and image file."""
0191 self.file.Write()
0192 self.page(close=True)
0193
0194
0195 class PythiaHistograms(Histograms):
0196 """PYHTIA-specific histograms"""
0197
0198 def __init__(self, file):
0199 super(PythiaHistograms, self).__init__(file)
0200 name = self.name
0201 self.histograms += [
0202 th2d(name('trueX'), '$x_{Bj}$', 25, -5., 0., 25, -5., 0.,
0203 select='log10(x):log10(trueX)', opt='colz'),
0204 th2d(name('trueQ2'), '$Q^{2}$', 25, -1., 4., 25, -1., 4.,
0205 select='log10(QSquared):log10(trueQ2)', opt='colz'),
0206 th2d(name('trueY'), '$y$', 25, -0.1, 1.1, 25, -0.1, 1.1,
0207 select='y:trueY', opt='colz'),
0208 th2d(name('trueW2'), '$W^{2}$', 25, 1., 6., 25, 1., 6.,
0209 select='log10(WSquared):log10(trueW2)', opt='colz')
0210 ]
0211
0212
0213 class DjangohHistograms(Histograms):
0214 """DJANGOH-specific histograms"""
0215
0216 def __init__(self, file):
0217 super(DjangohHistograms, self).__init__(file)
0218 name = self.name
0219 self.histograms += [
0220 th2d(name('dtrueX'), '$x_{Bj}$', 25, -5., 0., 25, -5., 0.,
0221 select='log10(x):log10(dtrueX)', opt='colz'),
0222 th2d(name('dtrueQ2'), '$Q^{2}$', 25, -1., 4., 25, -1., 4.,
0223 select='log10(QSquared):log10(dtrueQ2)', opt='colz'),
0224 th2d(name('dtrueY'), '$y$', 25, -0.1, 1.1, 25, -0.1, 1.1,
0225 select='y:dtrueY', opt='colz'),
0226 th2d(name('dtrueW2'), '$W^{2}$', 25, 1., 6., 25, 1., 6.,
0227 select='log10(WSquared):log10(dtrueW2)', opt='colz')
0228 ]
0229
0230
0231
0232 HISTOGRAMS = {'pythia': PythiaHistograms,
0233 'lepto': Histograms,
0234 'milou': Histograms,
0235 'djangoh': DjangohHistograms,
0236 'pepsi': Histograms,
0237 'gmc_trans': Histograms}
0238
0239 def event_type(file):
0240 """Return the event type for a ROOT file.
0241
0242 e.g. 'pythia', 'lepto'.
0243
0244 """
0245 file.EICTree.GetEntry(0)
0246 return file.EICTree.event.ClassName().lstrip('erhic::Event').lower()
0247
0248 if __name__ == '__main__':
0249 options = parse()
0250
0251 names = [build(name, options.outdir, options.nevents)
0252 for name in options.file]
0253
0254 root_files = [ROOT.TFile(name, 'update') for name in names]
0255
0256 histograms = [HISTOGRAMS[event_type(file)](file) for file in root_files]
0257
0258 for histogram in histograms:
0259 histogram.make()
0260 histogram.close()