Fork me on GitHub

source: svn/trunk/readers/DelphesProMC.py

Last change on this file was 1253, checked in by Pavel Demin, 11 years ago

remove import cStringIO from DelphesProMC.py

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