ATLAS Offline Software
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 
7 import AthenaCommon.SystemOfUnits as Units
8 from AthenaPython import PyAthena
9 StatusCode = PyAthena.StatusCode
10 import ROOT
11 
12 
13 def 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 
28 class 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 
291 class 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 
368 class 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
grepfile.info
info
Definition: grepfile.py:38
python.PyKernel.retrieve
def retrieve(aClass, aKey=None)
Definition: PyKernel.py:110
python.Lib.PyWriteThinnedData.sg
sg
storegate
Definition: AthenaExamples/AthExThinning/python/Lib.py:44
python.Lib.PyWriteThinnedData.test
def test(self, testName)
Definition: AthenaExamples/AthExThinning/python/Lib.py:63
python.Lib.PyWriteThinnedData.Particles
Particles
init base class
Definition: AthenaExamples/AthExThinning/python/Lib.py:36
SystemOfUnits
test_pyathena.px
px
Definition: test_pyathena.py:18
python.Lib.PyWriteThinnedData
Definition: AthenaExamples/AthExThinning/python/Lib.py:28
python.Lib.PyReadFatObject.sg
sg
Definition: AthenaExamples/AthExThinning/python/Lib.py:385
python.Lib.PyReadNonThinnedData
Definition: AthenaExamples/AthExThinning/python/Lib.py:291
python.Lib.PyWriteThinnedData.filter
filter
Datavector stuff.
Definition: AthenaExamples/AthExThinning/python/Lib.py:50
python.Lib.PyReadFatObject.__init__
def __init__(self, name="PyReadFatObject", **kw)
Definition: AthenaExamples/AthExThinning/python/Lib.py:370
python.Lib.PyReadFatObject.fatobject
fatobject
Definition: AthenaExamples/AthExThinning/python/Lib.py:377
PyAthena::Alg::initialize
virtual StatusCode initialize() override
Definition: PyAthenaAlg.cxx:60
PyAthena::Alg::execute
virtual StatusCode execute() override
Definition: PyAthenaAlg.cxx:93
python.Lib.PyReadFatObject.particles
particles
init base class
Definition: AthenaExamples/AthExThinning/python/Lib.py:376
python.Lib.keep
def keep(dec, l, op='set')
Definition: AthenaExamples/AthExThinning/python/Lib.py:13
python.Lib.PyWriteThinnedData.doThinningTest3
def doThinningTest3(self, iparticles, dec)
Definition: AthenaExamples/AthExThinning/python/Lib.py:243
PyAthena::Alg::finalize
virtual StatusCode finalize() override
Definition: PyAthenaAlg.cxx:86
python.Lib.PyReadNonThinnedData.sg
sg
storegate
Definition: AthenaExamples/AthExThinning/python/Lib.py:306
python.Lib.PyWriteThinnedData.Elephantino
Elephantino
Definition: AthenaExamples/AthExThinning/python/Lib.py:38
python.Lib.PyReadNonThinnedData.__init__
def __init__(self, name="PyReadNonThinnedData", **kw)
Definition: AthenaExamples/AthExThinning/python/Lib.py:294
python.Lib.PyWriteThinnedData.doThinningTest1
def doThinningTest1(self, particles, dec)
Definition: AthenaExamples/AthExThinning/python/Lib.py:159
plotBeamSpotVxVal.range
range
Definition: plotBeamSpotVxVal.py:195
python.Lib.PyReadFatObject
Definition: AthenaExamples/AthExThinning/python/Lib.py:368
python.Lib.PyReadNonThinnedData.Decay
Decay
Definition: AthenaExamples/AthExThinning/python/Lib.py:300
python.Lib.PyWriteThinnedData.Filter
Filter
Definition: AthenaExamples/AthExThinning/python/Lib.py:39
python.Lib.PyWriteThinnedData.doThinningTest2
def doThinningTest2(self, particles, dec)
Definition: AthenaExamples/AthExThinning/python/Lib.py:202
debug
const bool debug
Definition: MakeUncertaintyPlots.cxx:53
python.Lib.PyWriteThinnedData.Decay
Decay
Definition: AthenaExamples/AthExThinning/python/Lib.py:37
AthCommonMsg< Algorithm >::msg
MsgStream & msg() const
Definition: AthCommonMsg.h:24
PyAthena::Alg
Definition: PyAthenaAlg.h:33
python.Lib.PyReadNonThinnedData.Elephantino
Elephantino
Definition: AthenaExamples/AthExThinning/python/Lib.py:301
python.Lib.PyWriteThinnedData.__init__
def __init__(self, name="PyWriteThinnedData", **kw)
Definition: AthenaExamples/AthExThinning/python/Lib.py:31
python.Lib.PyReadNonThinnedData.Particles
Particles
init base class
Definition: AthenaExamples/AthExThinning/python/Lib.py:299
error
Definition: IImpactPoint3dEstimator.h:70
Generated on Thu Nov 7 2024 21:20:11 for ATLAS Offline Software by   doxygen 1.8.18