ATLAS Offline Software
Loading...
Searching...
No Matches
AthenaExamples/AthExThinning/python/Lib.py
Go to the documentation of this file.
1# Copyright (C) 2002-2023 CERN for the benefit of the ATLAS collaboration
2
3# @file: AthExThinning/python/Lib.py
4# @purpose: a set of Py-components to tests py-thinning
5# @author: Sebastien Binet <binet@cern.ch>
6
7import AthenaCommon.SystemOfUnits as Units
8from AthenaPython import PyAthena
9StatusCode = PyAthena.StatusCode
10import ROOT
11
12
13def keep (dec, l, op = 'set'):
14 if op == 'and':
15 opflag = 1
16 elif op == 'or':
17 opflag = 2
18 else:
19 opflag = 0
20 if isinstance (l, list):
21 for i, x in enumerate(l):
22 dec.keep (i, x, opflag)
23 else:
24 dec.keep (l, opflag)
25 return
26
27
28class PyWriteThinnedData( PyAthena.Alg ):
29 """A simple python algorithm to thin data
30 """
31 def __init__(self, name = "PyWriteThinnedData", **kw):
32
33 kw['name'] = name
34 super(PyWriteThinnedData,self).__init__(**kw)
35
36 self.Particles = kw.get('Particles', "Particles")
37 self.Decay = kw.get('Decay', "TwoBodyDecay")
38 self.Elephantino = kw.get('Elephantino', "PinkElephantino")
39 self.Filter = kw.get('Filter', [False]*10)
40
41 def initialize(self):
42 self.msg.info( "Initializing %s", self.name )
43
44 self.sg = PyAthena.py_svc("StoreGateSvc")
45
46
47 import cppyy
48 cppyy.load_library("libAthExThinningEventDict")
49 from RootUtils import PyROOTFixes # noqa: 401
50 self.filter = self.Filter
51 return StatusCode.Success
52
53 def execute(self):
54 self.msg.debug( "Executing %s...", self.name )
55 allGood = True
56 for t in ("test1","test2","test3",):
57 if self.test(t) != StatusCode.Success:
58 self.msg.warning( "Could not perform 'thinning %s'",t )
59 allGood = False
60 if allGood: return StatusCode.Success
61 return StatusCode.Failure
62
63 def test(self, testName):
64 self.msg.info( "Performing thinning test [%s]...", testName )
65
66 # fetch particles
67 particlesName = "%s_%s" % (self.Particles,testName)
68 particles = self.sg.retrieve( "AthExParticles", particlesName )
69 if particles is None:
70 self.msg.warning( "Could not fetch particles at [%s] !!",
71 self.Particles )
72 return StatusCode.Failure
73 self.msg.info( "particles: %i", particles.size() )
74
75 # fetch iparticles
76 iparticles = self.sg.retrieve( "AthExIParticles", particlesName )
77 if iparticles is None:
78 self.msg.warning( "Could not fetch iparticles at [%s] !!",
79 self.Particles )
80 return StatusCode.Failure
81 self.msg.info( "iparticles: %i", iparticles.size() )
82
83 if ( iparticles.size() != particles.size() or
84 iparticles.at(0).px() != particles.at(0).px() ):
85 self.msg.error( "symlinked containers are corrupted: " )
86 self.msg.error( " #iparticles: %i", iparticles.size() )
87 self.msg.error( " # particles: %i", particles.size() )
88 self.msg.error( " ipx[0] = %r", iparticles.at(0).px() )
89 self.msg.error( " px[0] = %r", particles.at(0).px() )
90 return StatusCode.Failure
91
92 # fetch decay
93 decay = self.sg.retrieve( "AthExDecay",
94 "%s_%s" % (self.Decay, testName) )
95 if decay is None:
96 self.msg.warning( "Could not fetch Decay at [%s] !!",
97 self.Decay )
98 return StatusCode.Failure
99
100 # fetch elephantino
101 elephantino = self.sg.retrieve( "AthExElephantino",
102 "%s_%s" % (self.Elephantino, testName) )
103 if elephantino is None:
104 self.msg.warning( "Could not fetch Elephantino at [%s] !!",
105 self.Elephantino )
106 return StatusCode.Failure
107
108 self.msg.info( "IN particles: %i", particles.size() )
109 self.msg.info( "IN decay:" )
110 self.msg.info( " p1: px= %r", decay.p1().px() / Units.GeV )
111 self.msg.info( " p2: px= %r", decay.p2().px() / Units.GeV )
112 self.msg.info( " l1: px= %r", decay.l1().px() / Units.GeV )
113 self.msg.info( " l2: px= %r", decay.l2().px() / Units.GeV )
114
115 self.msg.info( "IN elephantino:" )
116 self.msg.info( " leg1: px= %r", elephantino.leg1().px() / Units.GeV )
117 self.msg.info( " leg2: px= %r", elephantino.leg2().px() / Units.GeV )
118 self.msg.info( " leg3: px= %r", elephantino.leg3().px() / Units.GeV )
119 self.msg.info( " leg4: px= %r", elephantino.leg4().px() / Units.GeV )
120 self.msg.info( " ear1: px= %r", elephantino.ear1().px() / Units.GeV )
121 self.msg.info( " ear2: px= %r", elephantino.ear2().px() / Units.GeV )
122
123
124 dec = ROOT.SG.ThinningDecision (particlesName)
125 if testName == "test1":
126 if self.doThinningTest1( particles, dec) != StatusCode.Success:
127 self.msg.warning( "Could not exercize Thinning !!" )
128 elif testName == "test2":
129 if self.doThinningTest2( particles, dec) != StatusCode.Success:
130 self.msg.warning( "Could not exercize Thinning !!" )
131 elif testName == "test3":
132 if self.doThinningTest3(iparticles, dec) != StatusCode.Success:
133 self.msg.warning( "Could not exercize Thinning !!" )
134 else:
135 self.msg.error( "Unknown test [%s]", testName )
136 return StatusCode.Failure
137 dec.lock()
138 if not self.sg.record (dec, particlesName + '_THINNED_StreamUSR_0'):
139 return StatusCode.Failure
140 ROOT.SetOwnership (dec, False)
141
142 self.msg.info( "Decay is now:" )
143 self.msg.info( " p1: px= %r", decay.p1().px() / Units.GeV )
144 self.msg.info( " p2: px= %r", decay.p2().px() / Units.GeV )
145 self.msg.info( " l1: px= %r", decay.l1().px() / Units.GeV )
146 self.msg.info( " l2: px= %r", decay.l2().px() / Units.GeV )
147
148 self.msg.info( "Elephantino is now: " )
149 self.msg.info( " leg1: px= %r", elephantino.leg1().px() / Units.GeV )
150 self.msg.info( " leg2: px= %r", elephantino.leg2().px() / Units.GeV )
151 self.msg.info( " leg3: px= %r", elephantino.leg3().px() / Units.GeV )
152 self.msg.info( " leg4: px= %r", elephantino.leg4().px() / Units.GeV )
153 self.msg.info( " ear1: px= %r", elephantino.ear1().px() / Units.GeV )
154 self.msg.info( " ear2: px= %r", elephantino.ear2().px() / Units.GeV )
155
156 self.msg.info( "[%s] has been performed.", testName )
157 return StatusCode.Success
158
159 def doThinningTest1(self, particles, dec):
160 filter = self.filter[:]
161 RemovedIdx = ROOT.SG.ThinningDecisionBase.RemovedIdx
162 self.msg.info( "Particles | filter :" )
163 for i in range(particles.size()):
164 if filter[i]: kr = "keep"
165 else: kr = "remove"
166 self.msg.info( "%9s | %s", (i+1)*10, kr )
167 self.msg.info( "="*19 )
168
169 filter[len(filter)//2:] = [True]*(len(filter)//2)
170 self.msg.info( "Filter %r", filter )
171
172 self.msg.info( "... Processing [pre-thinning] ..." )
173 keep (dec, filter)
174
175 self.msg.info( "======== Index table =========" )
176 tmp = ROOT.SG.ThinningDecisionBase (dec)
177 tmp.buildIndexMap()
178 import os
179 os.sys.stdout.flush()
180 for i in range(particles.size()):
181 newIdx = tmp.index(i)
182 if newIdx == RemovedIdx: newIdx = "-"
183 self.msg.info( " idx %i -> %s", i, newIdx )
184
185 filter = self.filter[:]
186 filter[:len(filter)//2] = [True]*(len(filter)//2)
187 self.msg.info( "Filter %r", filter )
188
189 self.msg.info( "... Processing [thinning] ..." )
190 keep (dec, filter, 'and')
191
192 self.msg.info( "======== Index table =========" )
193 tmp = ROOT.SG.ThinningDecisionBase (dec)
194 tmp.buildIndexMap()
195 for i in range(particles.size()):
196 newIdx = tmp.index(i)
197 if newIdx == RemovedIdx: newIdx = "-"
198 self.msg.info( " idx %i -> %s", i, newIdx )
199
200 return StatusCode.Success
201
202 def doThinningTest2(self, particles, dec):
203 filter = self.filter[:]
204 RemovedIdx = ROOT.SG.ThinningDecisionBase.RemovedIdx
205 self.msg.info( "Particles | filter :" )
206 for i in range(particles.size()):
207 if filter[i]: kr = "keep"
208 else: kr = "remove"
209 self.msg.info( "%9s | %s", (i+1)*10, kr )
210 self.msg.info( "="*19 )
211
212 filter[len(filter)//2:] = [False]*(len(filter)//2)
213 self.msg.info( "Filter %s", filter )
214
215 self.msg.info( "... Processing [pre-thinning] ..." )
216 keep (dec, filter)
217
218 self.msg.info( "======== Index table =========" )
219 tmp = ROOT.SG.ThinningDecisionBase (dec)
220 tmp.buildIndexMap()
221 for i in range(particles.size()):
222 newIdx = tmp.index(i)
223 if newIdx == RemovedIdx: newIdx = "-"
224 self.msg.info( " idx %i -> %s", i, newIdx )
225
226 filter = self.filter[:]
227 filter[:len(filter)//2] = [False]*(len(filter)//2)
228 self.msg.info( "Filter %s", filter )
229
230 self.msg.info( "... Processing [thinning] ..." )
231 keep (dec, filter, 'or')
232
233 self.msg.info( "======== Index table =========" )
234 tmp = ROOT.SG.ThinningDecisionBase (dec)
235 tmp.buildIndexMap()
236 for i in range(particles.size()):
237 newIdx = tmp.index(i)
238 if newIdx == RemovedIdx: newIdx = "-"
239 self.msg.info( " idx %i -> %s", i, newIdx )
240
241 return StatusCode.Success
242
243 def doThinningTest3(self, iparticles, dec):
244 filter = self.filter[:]
245 RemovedIdx = ROOT.SG.ThinningDecisionBase.RemovedIdx
246 self.msg.info( "IParticles | filter :" )
247 for i in range(iparticles.size()):
248 if filter[i]: kr = "keep"
249 else: kr = "remove"
250 self.msg.info( "%9s | %s", (i+1)*10, kr )
251 self.msg.info( "="*19 )
252
253 filter[len(filter)//2:] = [True]*(len(filter)//2)
254 self.msg.info( "Filter %r", filter )
255
256 self.msg.info( "... Processing [pre-thinning] ..." )
257 keep (dec, filter)
258
259 self.msg.info( "======== Index table =========" )
260 tmp = ROOT.SG.ThinningDecisionBase (dec)
261 tmp.buildIndexMap()
262 for i in range(iparticles.size()):
263 newIdx = tmp.index(i)
264 if newIdx == RemovedIdx: newIdx = "-"
265 self.msg.info( " idx %i -> %s", i, newIdx )
266
267 filter = self.filter[:]
268 filter[:len(filter)//2] = [True]*(len(filter)//2)
269 self.msg.info( "Filter %r", filter )
270
271 self.msg.info( "... Processing [thinning] ..." )
272 keep (dec, filter, 'and')
273
274 self.msg.info( "======== Index table =========" )
275 tmp = ROOT.SG.ThinningDecisionBase (dec)
276 tmp.buildIndexMap()
277 for i in range(iparticles.size()):
278 newIdx = tmp.index(i)
279 if newIdx == RemovedIdx: newIdx = "-"
280 self.msg.info( " idx %i -> %s", i, newIdx )
281
282 return StatusCode.Success
283
284 def finalize(self):
285 self.msg.info( "Finalizing %s...", self.name )
286 return StatusCode.Success
287
288 pass # PyWriteThinnedData
289
290
291class PyReadNonThinnedData( PyAthena.Alg ):
292 """A simple python algorithm to read non-thinned data
293 """
294 def __init__(self, name = "PyReadNonThinnedData", **kw):
295
296 kw['name'] = name
297 super(PyReadNonThinnedData,self).__init__(**kw)
298
299 self.Particles = kw.get('Particles', "Particles")
300 self.Decay = kw.get('Decay', "TwoBodyDecay")
301 self.Elephantino = kw.get('Elephantino', "PinkElephantino")
302
303 def initialize(self):
304 self.msg.info( "Initializing %s", self.name )
305
306 self.sg = PyAthena.py_svc("StoreGateSvc")
307
308
309 import cppyy
310 cppyy.load_library("libAthExThinningEventDict")
311 from RootUtils import PyROOTFixes # noqa: F401
312 return StatusCode.Success
313
314 def execute(self):
315 sg = self.sg
316 _warning = self.msg.warning
317 _info = self.msg.info
318 _info( "Executing %s...", self.name )
319 for test in ("test1", "test2", "test3"):
320 key = "%s_%s"%(self.Particles,test)
321 particles = sg.retrieve("AthExParticles", key)
322 if particles is None:
323 _warning("Could not fetch particles at [%s] !!", key)
324 return StatusCode.Failure
325
326 iparticles = sg.retrieve("AthExIParticles", key)
327 if iparticles is None:
328 _warning("Could not fetch iparticles at [%s] !!", key)
329 return StatusCode.Failure
330
331 key = "%s_%s" % (self.Decay, test)
332 decay = sg.retrieve("AthExDecay", key)
333 if decay is None:
334 _warning("Could not fetch Decay at [%s] !!", key)
335 return StatusCode.Failure
336
337 key = "%s_%s" % (self.Elephantino, test)
338 elephantino = self.sg.retrieve("AthExElephantino", key)
339 if elephantino is None:
340 _warning("Could not fetch Elephantino at [%s] !!", key)
341 return StatusCode.Failure
342
343 _info( "Test: %s", test)
344 _info( "IN particles: %i", particles.size() )
345 _info( "IN iparticles: %i", iparticles.size() )
346 _info( "IN decay:" )
347 _info( " p1: px= %r", decay.p1().px() / Units.GeV )
348 _info( " p2: px= %r", decay.p2().px() / Units.GeV )
349 _info( " l1: px= %r", decay.l1().px() / Units.GeV )
350 _info( " l2: px= %r", decay.l2().px() / Units.GeV )
351
352 _info( "IN elephantino:" )
353 _info( " leg1: px= %r", elephantino.leg1().px() / Units.GeV )
354 _info( " leg2: px= %r", elephantino.leg2().px() / Units.GeV )
355 _info( " leg3: px= %r", elephantino.leg3().px() / Units.GeV )
356 _info( " leg4: px= %r", elephantino.leg4().px() / Units.GeV )
357 _info( " ear1: px= %r", elephantino.ear1().px() / Units.GeV )
358 _info( " ear2: px= %r", elephantino.ear2().px() / Units.GeV )
359 return StatusCode.Success
360
361 def finalize(self):
362 self.msg.info( "Finalizing %s...", self.name )
363 return StatusCode.Success
364
365 pass # PyReadNonThinnedData
366
367
368class PyReadFatObject (PyAthena.Alg):
369 """a simple algorithm to read back AthExFatObject"""
370 def __init__ (self, name="PyReadFatObject", **kw):
371
372 kw['name'] = name
373 super (PyReadFatObject, self).__init__ (**kw)
374
375
376 self.particles = kw.get ("particles", "AthExParticles")
377 self.fatobject = kw.get ("fatobject", "AthExFatObject")
378
379 return
380
381 def initialize(self):
382 _info = self.msg.info
383 _info('==> initialize...')
384 # retrieve storegate
385 self.sg = PyAthena.py_svc ("StoreGateSvc")
386 if not self.sg:
387 self.msg.error ("could not retrieve the event store !")
388 return StatusCode.Failure
389
390 _info ("input particles: [%s]", self.particles)
391 _info ("input fatobject: [%s]", self.fatobject)
392
393 return StatusCode.Success
394
395 def execute(self):
396 _info = self.msg.info
397 particles = self.sg[self.particles]
398 if not particles:
399 self.msg.error ("could not retrieve particles at [%s] !",
400 self.particles)
401 return StatusCode.Failure
402 _info ("particles' size: %i", particles.size())
403
404 fatobject = self.sg[self.fatobject]
405 if not fatobject:
406 self.msg.error ("could not retrieve fatobject at [%s] !",
407 self.fatobject)
408 return StatusCode.Failure
409
410 raw = fatobject.particle()
411 lnk = fatobject.particleLink()
412 _info ("fat.raw.px: %s",
413 "slimmed" if not raw else (raw.px()/Units.GeV))
414 _info ("fat.lnk.e : %s",
415 "slimmed" if not lnk.isValid() else (lnk.e() /Units.GeV))
416
417 return StatusCode.Success
418
419 def finalize(self):
420 self.msg.info('==> finalize...')
421 return StatusCode.Success
422
423 # class PyReadFatObject
debug
const bool debug
Definition MakeUncertaintyPlots.cxx:53
AthCommonMsg< Algorithm >::msg
MsgStream & msg() const
Definition AthCommonMsg.h:24
PyAthena::Alg::execute
virtual StatusCode execute() override
Definition PyAthenaAlg.cxx:93
PyAthena::Alg::finalize
virtual StatusCode finalize() override
Definition PyAthenaAlg.cxx:86
PyAthena::Alg::initialize
virtual StatusCode initialize() override
Definition PyAthenaAlg.cxx:60
python.Lib.PyReadFatObject.__init__
__init__(self, name="PyReadFatObject", **kw)
Definition AthenaExamples/AthExThinning/python/Lib.py:370
python.Lib.PyReadNonThinnedData.__init__
__init__(self, name="PyReadNonThinnedData", **kw)
Definition AthenaExamples/AthExThinning/python/Lib.py:294
python.Lib.PyWriteThinnedData.__init__
__init__(self, name="PyWriteThinnedData", **kw)
Definition AthenaExamples/AthExThinning/python/Lib.py:31
Generated on for ATLAS Offline Software by doxygen 1.14.0