eic-smear/eicpy/qaplots.py

0001 #!/usr/bin/env python
0002
0003 '''
0004 Suite of standard QA plots for EIC trees to check they are filled properly.
0005 It's fairly slow, but you don't need to run many events to check that all
0006 values are being filled properly.
0007 '''
0008
0009 import argparse
0010 from collections import namedtuple
0011
0012 # Configure the environment for ROOT/eic-smear.
0013 # Do this before importing ROOT, as initialise() will
0014 # attempt to set the environment to find libPyROOT
0015 # if it isn't already.
0016 import eicpy
0017 eicpy.initialise(False)
0018
0019 import ROOT
0020
0021 from binning import bin_log10
0022
0023 # Don't pass command line arguments to ROOT so it
0024 # doesn't intercept use of --help: we want the help
0025 # to be printed for this script, not ROOT.
0026 # root.cern.ch/phpBB3/viewtopic.php?f=14&t=13582
0027 ROOT.PyConfig.IgnoreCommandLineOptions = True
0028
0029 # Speed things up a bit
0030 ROOT.SetSignalPolicy(ROOT.kSignalFast)
0031
0032 # No need to draw graphics to screen
0033 ROOT.gROOT.SetBatch(True)
0034
0035 class LogAxis(namedtuple('LogAxis', 'x y z')):
0036     '''
0037     Utility for setting ROOT pad logarithmic axes.
0038     '''
0039     def apply(self, pad):
0040         '''
0041         Apply the log axis settings to a ROOT pad.
0042         '''
0043         pad.SetLogx(self.x != 0)
0044         pad.SetLogy(self.y != 0)
0045         pad.SetLogz(self.z != 0)
0046
0047     def rebin(self, hist):
0048         if self.x:
0049             bin_log10(hist.GetXaxis())
0050         if self.y and hist.GetDimension() > 1:
0051             bin_log10(hist.GetYaxis())
0052         if self.z and hist.GetDimension() > 2:
0053             bin_log10(hist.GetZaxis())
0054
0055
0056 class THWrapper(object):
0057     '''
0058     Simple wrapper around ROOT histogram adding some functionality.
0059     '''
0060     def __init__(self, h, log, option):
0061         '''
0062         Initialise with a histogram and a LogAxis giving whether
0063         each axis (x, y, z) should be log10 (1) or not (0).
0064         '''
0065         self.h = h
0066         self.log = log
0067         if not h.GetSumw2():
0068             self.h.Sumw2()
0069         self.h.SetOption(option)
0070         self.log.rebin(self.h)
0071
0072     def draw(self, pad):
0073         '''
0074         Draw the histogram to a ROOT pad.
0075         '''
0076         initial = ROOT.gPad
0077         pad.cd()
0078         self.log.apply(pad)
0079         self.h.Draw()
0080         ROOT.gPad = initial
0081
0082
0083 class EventHists(object):
0084     '''
0085     A collection of standard histograms for event-wise properties.
0086     '''
0087     def __init__(self):
0088         self.hists = {}
0089         self.hists['process'] = THWrapper(
0090             ROOT.TH1D('process', 'Event process number', 200, 0., 200.),
0091             LogAxis(0, 1, 0), '')
0092         self.hists['tracks'] = THWrapper(
0093             ROOT.TH1D('tracks', 'Number of tracks per event', 100, 0., 200),
0094             LogAxis(0, 0, 0), '')
0095 #        self.hists['eInNucl'] = THWrapper(
0096 #            ROOT.TH1D('eInNucl', 'E of incident lepton in nuclear rest frame',
0097 #                      40, 0., 10000),
0098 #            LogAxis(0, 0, 0), '')
0099 #        self.hists['eScNucl'] = THWrapper(
0100 #            ROOT.TH1D('eScNucl', 'E of scattered lepton in nuclear rest frame',
0101 #                      40, 0., 10000),
0102 #            LogAxis(0, 0, 0), '')
0103         self.hists['Q2x'] = THWrapper(
0104             ROOT.TH2D('Q2x', 'Q^{2} vs. x', 30, 1.e-5, 10., 16, 0.1, 1000.),
0105             LogAxis(1, 1, 1), 'colz')
0106         self.hists['y'] = THWrapper(
0107             ROOT.TH1D('y', 'y', 48, -0.1, 1.1), LogAxis(0, 1, 0), '')
0108         self.hists['W2'] = THWrapper(
0109             ROOT.TH1D('W2', 'W^{2} (GeV^{2})', 60, 0.1, 1.e5),
0110             LogAxis(1, 0, 0), '')
0111         self.hists['nu'] = THWrapper(
0112             ROOT.TH1D('nu', '#nu (GeV)', 60, 0.1, 1.e5), LogAxis(1, 0, 0), '')
0113         # Jacquet-Blondel and double-angle kinematics.
0114         # Clone the original histograms to maintain the same axis ranges.
0115         for i in ['y', 'Q2x', 'W2']:
0116             orig = self.hists[i]
0117             name = i + 'jb'
0118             self.hists[name] = THWrapper(orig.h.Clone(name), orig.log,
0119                                          orig.h.GetOption())
0120             self.hists[name].h.SetTitle(orig.h.GetTitle() +
0121                                         ' (Jacquet-Blondel)')
0122             name = i + 'da'
0123             self.hists[name] = THWrapper(orig.h.Clone(name), orig.log,
0124                                          orig.h.GetOption())
0125             self.hists[name].h.SetTitle(orig.h.GetTitle() +
0126                                         ' (double-angle)')
0127
0128     def output(self, pad, filename):
0129         '''
0130         Print all histograms to an output PostScript file.
0131         '''
0132         initial = ROOT.gPad
0133         pad.cd()
0134         for j in sorted(self.hists):
0135             self.hists[j].draw(pad)
0136             pad.Print(filename)
0137         ROOT.gPad = initial
0138
0139     def fill(self, event):
0140         '''
0141         '''
0142         self.hists['process'].h.Fill(event.GetProcess())
0143         self.hists['tracks'].h.Fill(event.GetNTracks())
0144 #        self.hists['eInNucl'].h.Fill(event.ELeptonInNucl)
0145 #        self.hists['eScNucl'].h.Fill(event.EScatteredInNucl)
0146         self.hists['Q2x'].h.Fill(event.GetX(), event.GetQ2())
0147         self.hists['y'].h.Fill(event.GetY())
0148         self.hists['W2'].h.Fill(event.GetW2())
0149         self.hists['nu'].h.Fill(event.GetNu())
0150         self.hists['Q2xjb'].h.Fill(event.GetXJacquetBlondel(),
0151                                    event.GetQ2JacquetBlondel())
0152         self.hists['yjb'].h.Fill(event.GetYJacquetBlondel())
0153         self.hists['W2jb'].h.Fill(event.GetW2JacquetBlondel())
0154         self.hists['Q2xda'].h.Fill(event.GetXDoubleAngle(),
0155                                    event.GetQ2DoubleAngle())
0156         self.hists['yda'].h.Fill(event.GetYDoubleAngle())
0157         self.hists['W2da'].h.Fill(event.GetW2DoubleAngle())
0158
0159
0160 class TrackHists(object):
0161     '''
0162     A collection of standard histograms for track-wise properties.
0163     Optionally specify a list of PDG codes to accept.
0164     Only tracks with one of those codes will be used when filling histograms.
0165     '''
0166     def __init__(self):
0167         # Particle status
0168         self.KS = THWrapper(
0169             ROOT.TH1D('KS', 'Track status code', 110, -10., 100.),
0170             LogAxis(0, 0, 0), '')
0171         # Particle parent index
0172         self.orig = THWrapper(
0173             ROOT.TH1D('orig', 'Track parent index', 210, -10., 200.),
0174             LogAxis(0, 0, 0), '')
0175         # Child indicies
0176         self.child = THWrapper(
0177             ROOT.TH2D('child', 'Child track indices;first;last',
0178                       210, -10., 200., 210, -10., 200.),
0179             LogAxis(0, 0, 0), 'colz')
0180         # y vs. x momentum
0181         self.pypx = THWrapper(
0182             ROOT.TH2D('pypx', 'p_{y} vs. p_{x} (GeV/c);p_{x};p_{y}',
0183                       60, -3., 3., 60, -3., 3.),
0184             LogAxis(0, 0, 1), 'colz')
0185         # Transverse vs. longitudinal momentum
0186         self.ptpz = THWrapper(
0187             ROOT.TH2D('ptpz', 'p_{T} vs. p_{z} (GeV/c);p_{z};p_{T}',
0188                       60, -40., 260., 60, -1., 11.),
0189             LogAxis(0, 0, 1), 'colz')
0190         # Energy vs. mass
0191         self.em = THWrapper(
0192             ROOT.TH2D('em', 'Energy vs. mass (final-state only);m;E',
0193                       60, -1., 11., 60, -40., 260.),
0194             LogAxis(0, 0, 1), 'colz')
0195         # Vertex y vs. x
0196         self.vxy = THWrapper(
0197             ROOT.TH2D('vxy', 'Track vertex y vs. x',
0198                       51, -100., 100., 51, -100., 100.),
0199             LogAxis(0, 0, 1), 'colz')
0200         # Energy fraction (z) vs. total momentum
0201         self.zp = THWrapper(
0202             ROOT.TH2D('zp', 'Energy fraction vs. p (final-state only);p;z',
0203                       60, -40., 260., 48, -0.1, 1.1),
0204             LogAxis(0, 0, 1), 'colz')
0205         # Angles theta vs. phi
0206         self.thetaphi = THWrapper(
0207             ROOT.TH2D('thetaphi', '#theta vs. #phi (rad)',
0208                       35, -0.4, 6.6, 40, -0.5, 3.5),
0209             LogAxis(0, 0, 1), 'colz')
0210         # Pseudorapidity vs. rapidity
0211         self.etay = THWrapper(
0212             ROOT.TH2D('etay', '#eta vs. rapidity',
0213                       40, -20., 20., 40, -20., 20.),
0214             LogAxis(0, 0, 1), 'colz')
0215         # p_T vs. polar angle in photon-proton frame
0216         self.gamma = THWrapper(
0217             ROOT.TH2D('gamma', 'p_{T} vs. #theta in #gamma-p frame',
0218                       40, -0.5, 3.5, 35, -1., 6.),
0219             LogAxis(0, 0, 1), 'colz')
0220 #        self.phi'] = THWrapper(
0221 #            ROOT.TH1D('phi', '#phi in #gamma-p frame', 35, -0.4, 6.6),
0222 #            LogAxis(0, 0, 0), '')
0223         # Feynman x
0224         self.xf = THWrapper(
0225             ROOT.TH1D('xf', 'Feynman x (final-state only)', 110, -1.1, 1.1),
0226             LogAxis(0, 0, 0), '')
0227
0228     def output(self, pad, filename):
0229         '''
0230         Print all histograms to an output PostScript file.
0231         '''
0232         initial = ROOT.gPad
0233         pad.cd()
0234         hists = [self.KS, self.child, self.em, self.etay, self.gamma, self.orig,
0235             self.ptpz, self.pypx, self.thetaphi, self.vxy, self.xf, self.zp]
0236         for j in hists:
0237             j.draw(pad)
0238             pad.Print(filename)
0239         ROOT.gPad = initial
0240
0241     def fill(self, tracks):
0242         '''
0243         Fill all histograms for a list of tracks
0244         '''
0245         # OK, this is a bit ugly looking, but localising all the track
0246         # and histogram methods outside the for loop saves quite a lot
0247         # of time when looping over millions of particles.
0248         # Here are all the Particle accessor methods we use:
0249         Particle = ROOT.erhic.ParticleMC
0250         status = Particle.GetStatus
0251         parent = Particle.GetParentIndex
0252         child1 = Particle.GetChild1Index
0253         childN = Particle.GetChildNIndex
0254         px = Particle.GetPx
0255         py = Particle.GetPy
0256         pz = Particle.GetPz
0257         pt = Particle.GetPt
0258         vert = Particle.GetVertex
0259         phi = Particle.GetPhi
0260         theta = Particle.GetTheta
0261         rap = Particle.GetRapidity
0262         eta = Particle.GetEta
0263         thetag = Particle.GetThetaVsGamma
0264         ptg = Particle.GetPtVsGamma
0265         e = Particle.GetE
0266         p = Particle.GetP
0267         m = Particle.GetM
0268         z = Particle.GetZ
0269         xf = Particle.GetXFeynman
0270         # Here are all the histogram Fill operations we need
0271         hks = self.KS.h.Fill
0272         horig = self.orig.h.Fill
0273         hchild = self.child.h.Fill
0274         hpypx = self.pypx.h.Fill
0275         hptpz = self.ptpz.h.Fill
0276         hvxy = self.vxy.h.Fill
0277         hthetaphi = self.thetaphi.h.Fill
0278         hetay = self.etay.h.Fill
0279         hgamma = self.gamma.h.Fill
0280         hem = self.em.h.Fill
0281         hzp = self.zp.h.Fill
0282         hxf = self.xf.h.Fill
0283         # We skip certain species for some of the histograms
0284         beams = [2212, 2112, 11]
0285         # Now we loop over the tracks
0286         for track in tracks:
0287             # Only fill final-state particles here as intermediate
0288             # state properties (e.g. negative mass for virtual photon)
0289             # can make it harder to spot issues.
0290             if track.GetStatus() != 1:
0291                 continue
0292             hks(status(track))
0293             horig(parent(track))
0294             hchild(child1(track), childN(track))
0295             hpypx(px(track), py(track))
0296             hptpz(pz(track), pt(track))
0297             vertex = vert(track)
0298             hvxy(vertex.x(), vertex.y())
0299             hthetaphi(phi(track), theta(track))
0300             hetay(rap(track), eta(track))
0301             hgamma(thetag(track), ptg(track))
0302             # Skip particles that are not in the final state, or may
0303             # be scattered beams, as some quantities don't make sense then.
0304             if track.GetPdgCode().Code() not in beams:
0305                 hem(m(track), e(track))
0306                 hzp(p(track), z(track))
0307                 hxf(xf(track))
0308
0309
0310 def print_progress(n, max):
0311     print 'Event', str(n).rjust(len(str(max))), 'of', max
0312
0313 def execute(tree, outname, maxevents):
0314     '''
0315     '''
0316     eicpy.root_logon()
0317     outfile = ROOT.TFile(outname.replace('.ps', '.root'), 'recreate')
0318     h = EventHists()
0319     t = TrackHists()
0320     nevents = min(maxevents, tree.GetEntries())
0321     for i in xrange(0, nevents):
0322         tree.GetEntry(i)
0323         h.fill(tree.event)
0324         tracks = [track for track in tree.event.GetTracks()]
0325         t.fill(tracks)
0326         if (i + 1) % (nevents / 10) == 0:
0327             print_progress(i + 1, nevents)
0328     window = ROOT.TCanvas()
0329     window.Print(''.join([outname, '[']))
0330     h.output(window, outname)
0331     t.output(window, outname)
0332     window.Print(''.join([outname, ']']))
0333     outfile.Write()
0334
0335 if __name__ == '__main__':
0336     # Process the command line arguments
0337     parser = argparse.ArgumentParser()
0338     parser.add_argument('infile', help='Input ROOT file')
0339     parser.add_argument('outfile', help='Output PostScript file')
0340     parser.add_argument('--nevents', type=int, default=1000000000, help='Maximum number of events')
0341     args = parser.parse_args()
0342     file = ROOT.TFile(args.infile, 'read')
0343     execute(file.EICTree, args.outfile, args.nevents)