Warning, /eic-smear/eicpy/tppmc2smear is written in an unsupported language. File is not indexed.
0001 #!/usr/bin/env python
0002
0003 '''
0004 Bridging script between TPPMC output and EIC smearing code.
0005
0006 You need both TPPMC and eic-smear installed, and the libraries to
0007 be accessible via (DY)LD_LIBRARY_PATH.
0008 ROOT should be installed with Python support (which it is by default).
0009
0010 First, it generates a ROOT tree in the EIC event format, specifically
0011 events of the class erhic::hadronic::EventPythia, from the TPPMC events.
0012 Only TPPMC hadrons are propagated to the EIC tree file, as the partons
0013 are not involved in smearing.
0014 Then, smearing is performed on the EIC tree.
0015 Users can define their own smearing performance either by
0016 1) modifying the subroutine build_detector() in tppmc2smear.py
0017 2) providing an external ROOT script. The external script should provide
0018 a function with the following name and signature:
0019 void define_smearing(Smear::Detector&);
0020 The function should define the Smear::Detector that guides smearing
0021 according to their needs.
0022 By default the intermediate EIC event tree will be saved, with the name
0023 unsmeared.<name>.root for a smeared file named <name>.root.
0024
0025 Run tppmc2smear.py --help for an explanation of all command line options.
0026
0027 Examples:
0028
0029 1) Create EIC/smeared trees:
0030 tppmc2smear.py tppmc.root myfile.root
0031 Generates myfile.root and unsmeared.myfile.root
0032
0033 2) Use smearing defined in an external script:
0034 tppmc2smear.py tppmc.root myfile.root --script detector.C
0035
0036 3) Delete intermediate unsmeared EIC file
0037 tppmc2smear.py tppmc.root myfile.root --clean
0038 '''
0039
0040 import argparse
0041 import math
0042 import os
0043 import subprocess
0044 import sys
0045
0046 import ROOT
0047
0048 # Don't pass command line arguments to ROOT so it
0049 # doesn't intercept use of --help: we want the help
0050 # to be printed for this script, not ROOT.
0051 # root.cern.ch/phpBB3/viewtopic.php?f=14&t=13582
0052 ROOT.PyConfig.IgnoreCommandLineOptions = True
0053
0054 # Speed things up a bit
0055 ROOT.SetSignalPolicy(ROOT.kSignalFast)
0056
0057 def open_tppmc_file(filename):
0058 '''
0059 Open the named ROOT file.
0060 Return a tuple of the file and the tree in the file.
0061 '''
0062 try:
0063 file = ROOT.TFile(filename, 'read')
0064 tree = file.Get('events')
0065 print "Opened file '{}' containing TTree '{}'".format(
0066 file.GetName(), tree.GetName())
0067 return file, tree
0068 except Exception as e:
0069 print 'error in open_tppmc_file', e
0070
0071 class Pdg:
0072 '''
0073 Class for look-up of static particle properties.
0074 Caches ROOT.TParticlePDG objects to avoid repeated
0075 calls to ROOT.TDatabasePDG.GetParticle().
0076 '''
0077 particles = {}
0078 database = ROOT.TDatabasePDG.Instance()
0079 pow = math.pow
0080 @classmethod
0081 def mass2(cls, pdgcode):
0082 # Populate the dict if this PDG code is not present
0083 if pdgcode not in cls.particles:
0084 cls.particles[pdgcode] = cls.database.GetParticle(pdgcode)
0085 # Return the mass if the particle is known, 0 if not
0086 if cls.particles[pdgcode]:
0087 return cls.pow(cls.particles[pdgcode].Mass(), 2.)
0088 else:
0089 return 0.
0090
0091
0092 '''
0093 def get_mass(pdg):
0094 Returns the mass (GeV/c^2) of the particle species identified
0095 by the PDG code, or zero if the particle cannot be determined.
0096 if ROOT.TDatabasePDG.Instance().GetParticle(pdg):
0097 return ROOT.TDatabasePDG.Instance().GetParticle(pdg).Mass()
0098 return 0.
0099 '''
0100 def tppmc_particle(tp, cme):
0101 '''
0102 Generate an erhic::hadronic::ParticleMC from a tppmc::Particle
0103 '''
0104 # TPPMC doesn't store particle energy, so
0105 # compute it from momentum and mass
0106 pdg = tp.Type()
0107 mom3 = tp.PxPyPz()
0108 energy = math.sqrt(mom3.Mag2() + math.pow(Pdg.mass(pdg), 2.))
0109 # energy = math.sqrt(mom3.Mag2() + math.pow(get_mass(pdg), 2.))
0110 momentum = ROOT.TLorentzVector(mom3, energy)
0111 # All TPPMC particles are final-state (status = 1).
0112 # Vertex and parent index is not recorded, so use 0.
0113 ep = ROOT.erhic.hadronic.ParticleMC(
0114 momentum, ROOT.TVector3(0., 0., 0.), pdg, 1, 0)
0115 ep.SetXFeynman(ep.GetPz() * 2. / cme)
0116 return ep
0117
0118 def qa(options):
0119 '''
0120 Make some simple plots comparing values in the TPPMC and
0121 EIC trees, to ensure values were properly copied.
0122 '''
0123 print type(options.tppmc)
0124 tppmcfile = ROOT.TFile(options.tppmc)
0125 tppmctree = tppmcfile.events
0126 tppmctree.AddFriend('EICTree', options.smear)
0127 window = ROOT.TCanvas()
0128 window.Print(options.qa + '[')
0129 option = 'box'
0130 for pattern in [
0131 'mHadrons.pz:mTracks.pz', 'mHadrons.PxPyPz().Eta():mTracks.eta',
0132 'mHadrons.type:mTracks.id', 'QSquared():EICTree.QSquared',
0133 'PartonX1():EICTree.x1', 'PartonX2():EICTree.x2']:
0134 tppmctree.Draw(pattern, '', option)
0135 window.Print(options.qa)
0136 tppmctree.Draw('EICTree.mTracks.xFeynman')
0137 window.Print(options.qa + ')')
0138
0139 def build_detector():
0140 '''
0141 An example of a possible EIC detector.
0142 Put your own smearing definitions here.
0143 '''
0144 emcal = ROOT.Smear.Device(ROOT.Smear.kE, '0.15*sqrt(E)',
0145 ROOT.Smear.kElectromagnetic)
0146 momentum = ROOT.Smear.Device(ROOT.Smear.kP, '0.005*P + 0.004*P*P')
0147 theta = ROOT.Smear.Device(ROOT.Smear.kTheta, '0.')
0148 phi = ROOT.Smear.Device(ROOT.Smear.kPhi, '0')
0149 pid = ROOT.Smear.ParticleID('PerfectPIDMatrix.dat')
0150
0151 detector = ROOT.Smear.Detector()
0152 detector.SetEventKinematicsCalculator('NM JB DA')
0153 for d in [emcal, momentum, theta, phi, pid]:
0154 detector.AddDevice(d)
0155
0156 return detector
0157
0158 def runsmear(filename, detector, outname = ''):
0159 '''
0160 Use the SmearTree routine to run smearing.
0161 '''
0162 out = outname
0163 if len(out) == 0:
0164 base, ext = os.path.splitext(os.path.basename(filename))
0165 out = '.'.join([base, 'smear.root'])
0166 ROOT.SmearTree(detector, filename, out)
0167
0168 def load_external_root_script(options):
0169 '''
0170 Check for the external ROOT script named by the command
0171 line options and load it into ROOT. Ensure that it provides
0172 the expected defined_smearing() function.
0173 Return True if everything is OK, or raise an exception
0174 if an error is encountered.
0175 '''
0176 if not os.path.exists(options.script):
0177 raise IOError(
0178 'External script {} does not exist'.format(options.script))
0179 ROOT.gROOT.LoadMacro(options.script)
0180 # Check that the user provided a function called define_smearing()
0181 # in their external script.
0182 if not ROOT.gROOT.GetGlobalFunction('define_smearing'):
0183 raise RuntimeError(
0184 '{} does not provide define_smearing()'.format(options.script))
0185 return True
0186
0187 def get_detector(options):
0188 '''
0189 Returns a ROOT.Smear.Detector(), either defined in this
0190 script or provided by an external ROOT script, depening
0191 on command line options.
0192 '''
0193 # Check for an external script
0194 if options.script and load_external_root_script(options):
0195 print 'Reading smearing from', options.script
0196 detector = ROOT.Smear.Detector()
0197 # The external script loaded a function called define_smearing().
0198 ROOT.define_smearing(detector)
0199 else:
0200 print 'Using tppmc2smear smearing definition'
0201 detector = build_detector()
0202 return detector
0203
0204 def generate(options):
0205 file, tree = open_tppmc_file(options.tppmc)
0206 eicfile = ROOT.TFile('.'.join(['unsmeared', options.smear]), 'recreate')
0207 eictree = ROOT.TTree('EICTree', 'Converted TPPMC event tree')
0208 eictree.Branch('event', ROOT.erhic.hadronic.EventPythiaPP())
0209 entries = tree.GetEntries()
0210 interval = entries / 10
0211 print 'Converting tree to suitable format...'
0212 event_loop(tree, eictree, entries, interval)
0213 eicfile.Write()
0214 # Now we have a tree corresponding to the input TPPMC tree, but
0215 # in a format that the smearing code will accept.
0216 runsmear(eicfile.GetName(), get_detector(options), options.smear)
0217 # Delete intermediate EIC file if requested
0218 if options.clean:
0219 subprocess.call(['rm', '-f', eicfile.GetName()])
0220
0221 class TreeReader(object):
0222 '''
0223 Generator class for entry-by-entry reading of ROOT TTree.
0224 Each call returns the next event in the tree.
0225 '''
0226 def __init__(self, tree):
0227 self.tree = tree
0228 self.entries = tree.GetEntries()
0229 self.get = ROOT.TTree.GetEntry
0230 def __call__(self):
0231 for i in xrange(0, self.entries):
0232 self.get(self.tree, i)
0233 yield self.tree.event
0234
0235 def event_loop(tree, eictree, entries, interval):
0236 # Not creating the vectors each time in the inner loop saves some time.
0237 vertex = ROOT.TVector3(0., 0., 0.)
0238 momentum = ROOT.TLorentzVector(0., 0., 0., 0.)
0239 # Read input tree event by event and copy properties to output event.
0240 read_events = TreeReader(tree)
0241 for te in read_events():
0242 event = ROOT.erhic.hadronic.EventPythiaPP(
0243 te.QSquared(), te.PartonX1(), te.PartonX2())
0244 # Convert and add hadrons to the eRHIC event
0245 hadrons = te.Hadrons()
0246 # This runs faster without having the particle creation
0247 # in a separate function.
0248 for i in hadrons:
0249 # TPPMC doesn't store particle energy, so
0250 # compute it from momentum and mass
0251 pdg = i.Type()
0252 mom3 = i.PxPyPz()
0253 energy = math.sqrt(mom3.Mag2() + Pdg.mass2(pdg))
0254 momentum.SetVect(mom3)
0255 momentum.SetE(energy)
0256 # All TPPMC particles are final-state (status = 1).
0257 # Vertex and parent index is not recorded, so use 0.
0258 ep = ROOT.erhic.hadronic.ParticleMC(momentum, vertex, pdg, 1, 0)
0259 ep.SetXFeynman(ep.GetPz() * 2. / te.mCentreOfMassEnergy)
0260 event.Add(ep)
0261 eictree.SetBranchAddress('event', event)
0262 eictree.Fill()
0263 # Print the number of events processed.
0264 # Pad the number correctly for nice formatting.
0265 if eictree.GetEntries() % interval == 0 and eictree.GetEntries() > 0:
0266 print str(eictree.GetEntries()).rjust(len(str(entries))), \
0267 'of', entries, 'events'
0268
0269 def process_arguments():
0270 '''
0271 Processes command line arguments.
0272 Returns an argparse.Namespace with the following fields:
0273 tppmc
0274 erhic
0275 qa
0276 script
0277 clean
0278 '''
0279 parser = argparse.ArgumentParser()
0280 parser.add_argument('tppmc', help='name of input TPPMC ROOT file')
0281 parser.add_argument('smear', help='name of output smeared ROOT file')
0282 parser.add_argument('--script', action='store',
0283 help='name of ROOT script defining smearing')
0284 parser.add_argument('--clean', action='store_const', const=True,
0285 default=False, help='delete intermediate EIC file')
0286 # This is just for testing, so suppress help
0287 parser.add_argument('--qa', action='store', default=None, help=argparse.SUPPRESS)
0288 options = parser.parse_args()
0289 return options
0290
0291 def execute():
0292 ROOT.gSystem.Load('libtppmc')
0293 ROOT.gSystem.Load('libeicsmear')
0294 options = process_arguments()
0295 if options.qa:
0296 qa(options)
0297 else:
0298 generate(options)
0299
0300 if __name__ == "__main__":
0301 execute()
