eic-smear/eicpy/verify.py

0001 #!/usr/bin/env python
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': ''} # Defaults
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 # Add as a new member function to TH1
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         # erhic::EventDis variables
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         # x, Q2, W2 and y, for each calculation method
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         # erhic::EventMC variables:
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 # Define Histograms to make for each generator
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     # Build ROOT files
0251     names = [build(name, options.outdir, options.nevents)
0252              for name in options.file]
0253     # Open ROOT files for updating with histograms
0254     root_files = [ROOT.TFile(name, 'update') for name in names]
0255     # Initialise histograms
0256     histograms = [HISTOGRAMS[event_type(file)](file) for file in root_files]
0257     # Fill, draw and write histograms
0258     for histogram in histograms:
0259         histogram.make()
0260         histogram.close()