[1252] | 1 | #!/usr/bin/env python
|
---|
| 2 |
|
---|
[1251] | 3 | import sys
|
---|
| 4 |
|
---|
| 5 | import urllib2
|
---|
| 6 | import cStringIO
|
---|
| 7 | import zipfile
|
---|
| 8 |
|
---|
| 9 | import ROOT
|
---|
| 10 |
|
---|
| 11 | import ProMCHeader_pb2
|
---|
| 12 | import ProMC_pb2
|
---|
| 13 |
|
---|
| 14 | ################################################################################
|
---|
| 15 | # HttpFile class is written by Eli Carter (retracile)
|
---|
| 16 | # http://stackoverflow.com/a/7852229
|
---|
| 17 |
|
---|
| 18 | class HttpFile(object):
|
---|
| 19 | def __init__(self, url):
|
---|
| 20 | self.url = url
|
---|
| 21 | self.offset = 0
|
---|
| 22 | self.length = -1
|
---|
| 23 |
|
---|
| 24 | def size(self):
|
---|
| 25 | if self.length < 0:
|
---|
| 26 | f = urllib2.urlopen(self.url)
|
---|
| 27 | self.length = int(f.headers["Content-length"])
|
---|
| 28 | return self.length
|
---|
| 29 |
|
---|
| 30 | def read(self, count = -1):
|
---|
| 31 | req = urllib2.Request(self.url)
|
---|
| 32 | if count < 0:
|
---|
| 33 | end = self.size() - 1
|
---|
| 34 | else:
|
---|
| 35 | end = self.offset + count - 1
|
---|
| 36 | req.headers['Range'] = "bytes=%s-%s" % (self.offset, end)
|
---|
| 37 | f = urllib2.urlopen(req)
|
---|
| 38 | data = f.read()
|
---|
| 39 | # FIXME: should check that we got the range expected, etc.
|
---|
| 40 | chunk = len(data)
|
---|
| 41 | if count >= 0:
|
---|
| 42 | assert chunk == count
|
---|
| 43 | self.offset += chunk
|
---|
| 44 | return data
|
---|
| 45 |
|
---|
| 46 | def seek(self, offset, whence = 0):
|
---|
| 47 | if whence == 0:
|
---|
| 48 | self.offset = offset
|
---|
| 49 | elif whence == 1:
|
---|
| 50 | self.offset += offset
|
---|
| 51 | elif whence == 2:
|
---|
| 52 | self.offset = self.size() + offset
|
---|
| 53 | else:
|
---|
| 54 | raise Exception("Invalid whence")
|
---|
| 55 |
|
---|
| 56 | def tell(self):
|
---|
| 57 | return self.offset
|
---|
| 58 |
|
---|
| 59 | ################################################################################
|
---|
| 60 |
|
---|
| 61 | def ConvertInput(name, momentumUnit, lengthUnit, branch, factory,
|
---|
| 62 | allParticleOutputArray, stableParticleOutputArray, partonOutputArray):
|
---|
| 63 |
|
---|
| 64 | pdg = ROOT.TDatabasePDG.Instance()
|
---|
| 65 |
|
---|
| 66 | data = ProMC_pb2.ProMCEvent()
|
---|
| 67 | data.ParseFromString(zip.read(name))
|
---|
| 68 |
|
---|
| 69 | event = data.event
|
---|
| 70 | particles = data.particles
|
---|
| 71 |
|
---|
| 72 | element = branch.NewEntry()
|
---|
| 73 |
|
---|
| 74 | element.Number = event.Number[0] if event.Number else 0
|
---|
| 75 | element.ProcessID = event.Process_ID[0] if event.Process_ID else 0
|
---|
| 76 | element.MPI = event.MPI[0] if event.MPI else 0
|
---|
| 77 | # element.Weight = event.Weight[0] if event.Weight else 0
|
---|
| 78 | element.Scale = event.Scale[0] if event.Scale else 0
|
---|
| 79 | element.AlphaQED = event.Alpha_QED[0] if event.Alpha_QED else 0
|
---|
| 80 | element.AlphaQCD = event.Alpha_QCD[0] if event.Alpha_QCD else 0
|
---|
| 81 | element.ID1 = event.ID1[0] if event.ID1 else 0
|
---|
| 82 | element.ID2 = event.ID2[0] if event.ID2 else 0
|
---|
| 83 | element.X1 = event.X1[0] if event.X1 else 0
|
---|
| 84 | element.X2 = event.X2[0] if event.X2 else 0
|
---|
| 85 | element.ScalePDF = event.Scale_PDF[0] if event.Scale_PDF else 0
|
---|
| 86 | element.PDF1 = event.PDF1[0] if event.PDF1 else 0
|
---|
| 87 | element.PDF2 = event.PDF2[0] if event.PDF2 else 0
|
---|
| 88 |
|
---|
| 89 | for i in range(len(particles.pdg_id)):
|
---|
| 90 | pid = particles.pdg_id[i]
|
---|
| 91 | status = particles.status[i]
|
---|
| 92 | px = particles.Px[i]
|
---|
| 93 | py = particles.Py[i]
|
---|
| 94 | pz = particles.Pz[i]
|
---|
| 95 | mass = particles.mass[i]
|
---|
| 96 |
|
---|
| 97 | x = particles.X[i]
|
---|
| 98 | y = particles.Y[i]
|
---|
| 99 | z = particles.Z[i]
|
---|
| 100 | t = particles.T[i]
|
---|
| 101 |
|
---|
| 102 | candidate = factory.NewCandidate()
|
---|
| 103 |
|
---|
| 104 | candidate.PID = pid
|
---|
| 105 | pdgCode = ROOT.TMath.Abs(candidate.PID)
|
---|
| 106 |
|
---|
| 107 | candidate.Status = status
|
---|
| 108 |
|
---|
[1252] | 109 | candidate.M1 = particles.mother1[i] if particles.mother1[i] < sys.maxint else -1
|
---|
| 110 | candidate.M2 = particles.mother2[i] if particles.mother2[i] < sys.maxint else -1
|
---|
[1251] | 111 |
|
---|
[1252] | 112 | candidate.D1 = particles.daughter1[i] if particles.daughter1[i] < sys.maxint else -1
|
---|
| 113 | candidate.D2 = particles.daughter2[i] if particles.daughter2[i] < sys.maxint else -1
|
---|
[1251] | 114 |
|
---|
| 115 | candidate.Mass = mass
|
---|
| 116 |
|
---|
| 117 | candidate.Momentum.SetXYZM(px, py, pz, mass)
|
---|
| 118 |
|
---|
| 119 | candidate.Position.SetXYZT(x, y, z, t)
|
---|
| 120 |
|
---|
| 121 | pdgParticle = pdg.GetParticle(pid)
|
---|
| 122 | if pdgParticle == None:
|
---|
| 123 | candidate.Charge = -999
|
---|
| 124 | continue
|
---|
| 125 | else:
|
---|
| 126 | candidate.Charge = int(pdgParticle.Charge()/3.0)
|
---|
| 127 |
|
---|
| 128 | allParticleOutputArray.Add(candidate)
|
---|
| 129 |
|
---|
| 130 | if status == 1:
|
---|
| 131 | stableParticleOutputArray.Add(candidate)
|
---|
| 132 | elif pdgCode <= 5 or pdgCode == 21 or pdgCode == 15:
|
---|
| 133 | partonOutputArray.Add(candidate)
|
---|
| 134 |
|
---|
| 135 | ################################################################################
|
---|
| 136 |
|
---|
| 137 | if len(sys.argv) < 2:
|
---|
| 138 | print " Usage: DelphesProMC.py config_file output_file input_file(s)"
|
---|
| 139 | sys.exit(1)
|
---|
| 140 |
|
---|
| 141 | ROOT.gSystem.Load('libDelphes')
|
---|
| 142 |
|
---|
| 143 | outputFile = ROOT.TFile(sys.argv[2], "CREATE")
|
---|
| 144 |
|
---|
| 145 | if not outputFile.IsOpen():
|
---|
| 146 | print "** ERROR: can't open", sys.argv[2]
|
---|
| 147 | sys.exit(1)
|
---|
| 148 |
|
---|
| 149 | treeWriter = ROOT.ExRootTreeWriter(outputFile, "Delphes")
|
---|
| 150 |
|
---|
| 151 | branchEvent = treeWriter.NewBranch("Event", ROOT.HepMCEvent.Class())
|
---|
| 152 |
|
---|
| 153 | confReader = ROOT.ExRootConfReader()
|
---|
| 154 | confReader.ReadFile(sys.argv[1])
|
---|
| 155 |
|
---|
| 156 | modularDelphes = ROOT.Delphes("Delphes")
|
---|
| 157 | modularDelphes.SetConfReader(confReader)
|
---|
| 158 | modularDelphes.SetTreeWriter(treeWriter)
|
---|
| 159 |
|
---|
| 160 | factory = modularDelphes.GetFactory()
|
---|
| 161 | allParticleOutputArray = modularDelphes.ExportArray("allParticles")
|
---|
| 162 | stableParticleOutputArray = modularDelphes.ExportArray("stableParticles")
|
---|
| 163 | partonOutputArray = modularDelphes.ExportArray("partons")
|
---|
| 164 |
|
---|
| 165 | modularDelphes.InitTask()
|
---|
| 166 |
|
---|
| 167 | for fileName in sys.argv[3:]:
|
---|
| 168 | print "** Reading", fileName
|
---|
| 169 |
|
---|
| 170 | if fileName.startswith("http://"):
|
---|
| 171 | file = HttpFile(fileName)
|
---|
| 172 | else:
|
---|
| 173 | file = open(fileName)
|
---|
| 174 |
|
---|
| 175 | zip = zipfile.ZipFile(file)
|
---|
| 176 |
|
---|
| 177 | numberOfEvents = len(zip.namelist())
|
---|
| 178 | if numberOfEvents <= 0: continue
|
---|
| 179 |
|
---|
| 180 | progressBar = ROOT.ExRootProgressBar(numberOfEvents - 1)
|
---|
| 181 |
|
---|
| 182 | # retrive information from the header file
|
---|
| 183 | header = ProMCHeader_pb2.ProMCHeader()
|
---|
| 184 | header.ParseFromString(zip.read("header"))
|
---|
| 185 | momentumUnit = float(header.MomentumUnit)
|
---|
| 186 | lengthUnit = float(header.LengthUnit)
|
---|
| 187 |
|
---|
| 188 | modularDelphes.Clear()
|
---|
| 189 | treeWriter.Clear()
|
---|
| 190 | eventCounter = 0
|
---|
| 191 | for name in zip.namelist():
|
---|
| 192 | eventCounter += 1
|
---|
| 193 | if not name.isdigit(): continue
|
---|
| 194 |
|
---|
| 195 | ConvertInput(name, momentumUnit, lengthUnit, branchEvent, factory,
|
---|
| 196 | allParticleOutputArray, stableParticleOutputArray, partonOutputArray)
|
---|
| 197 |
|
---|
| 198 | modularDelphes.ProcessTask()
|
---|
| 199 |
|
---|
| 200 | treeWriter.Fill()
|
---|
| 201 |
|
---|
| 202 | modularDelphes.Clear()
|
---|
| 203 | treeWriter.Clear()
|
---|
| 204 | progressBar.Update(eventCounter)
|
---|
| 205 |
|
---|
| 206 | progressBar.Update(eventCounter, eventCounter, ROOT.kTRUE)
|
---|
| 207 | progressBar.Finish()
|
---|
| 208 | zip.close()
|
---|
| 209 |
|
---|
| 210 | modularDelphes.FinishTask()
|
---|
| 211 | treeWriter.Write()
|
---|
| 212 |
|
---|
| 213 | del treeWriter
|
---|