Belle II Software light-2406-ragdoll
particlelistmerging.py
1
8import basf2
9import math
10import random
11import modularAnalysis as ma
12from ROOT import Belle2
13
14# disable doxygen check for this test
15# @cond
16
17
18class Generator(basf2.Module):
19 """Generate list of 10 electrons which have random momenta and one
20 electron where all momentum components are nan"""
21
22 def initialize(self):
23 """We need to register the mc particles"""
24
25 self.mcp = Belle2.PyStoreArray("MCParticles")
26 self.mcp.registerInDataStore()
27
28 def event(self):
29 """And then we generate particles"""
30 for i in range(10):
31 p = self.mcp.appendNew()
32 p.setPDG(11)
33 p.setMassFromPDG()
34 p.setMomentum(random.randrange(0, 5), random.randrange(0, 5), random.randrange(0, 5))
35
36 p = self.mcp.appendNew()
37 p.setPDG(11)
38 p.setMassFromPDG()
39 p.setMomentum(math.nan, math.nan, math.nan)
40
41
42class MergingChecker(basf2.Module):
43 """Check if merging works correctly"""
44
45 def initialize(self):
46 """Create particle list objects"""
47
48 self.all = Belle2.PyStoreObj('e-:allMC')
49 self.below2 = Belle2.PyStoreObj('e-:below2')
50 self.above3 = Belle2.PyStoreObj('e-:above3')
51 self.onlyall = Belle2.PyStoreObj('e-:onlyall')
52 self.goodmerge = Belle2.PyStoreObj('e-:goodmerge')
53 self.best = Belle2.PyStoreObj('e-:best')
54
55 def event(self):
56 """Check if the particle lists have the expected size"""
57 allsize = self.all.getListSize()
58 below2size = self.below2.getListSize()
59 above3size = self.above3.getListSize()
60 onlyallsize = self.onlyall.getListSize()
61 goodmergesize = self.goodmerge.getListSize()
62 bestsize = self.best.getListSize()
63 # The test fails if the number of particles in the merged lists differ from the sum of the desired input lists
64 assert onlyallsize == allsize
65 assert goodmergesize == below2size + above3size
66 assert bestsize == below2size + above3size
67
68
69# we create 10 events with 11 electrons each
70path = basf2.Path()
71path.add_module("EventInfoSetter", evtNumList=10)
72path.add_module(Generator())
73# load all electrons and add extraInfo so that they can be identified
74ma.fillParticleListFromMC("e-:allMC", "", path=path)
75ma.variablesToExtraInfo('e-:allMC', {'formula(0)': 'identifier', 'random': 'quality'}, path=path)
76# load electrons with px below 2
77ma.fillParticleListFromMC("e-:below2", "px < 2", path=path)
78ma.variablesToExtraInfo('e-:below2', {'formula(2)': 'identifier', 'formula(1 + random)': 'quality'}, option=2, path=path)
79# load electrons with px above 3
80ma.fillParticleListFromMC("e-:above3", "px > 3", path=path)
81ma.variablesToExtraInfo('e-:above3', {'formula(3)': 'identifier', 'formula(2 + random)': 'quality'}, option=2, path=path)
82
83# In the momentum range 0 < px < 2 the electrons from the list e-:below2 have the highest quality value
84# Electrons with a x-momentum between 2 and 3 are only part of the e-:all list.
85# For electrons with x-momenta above 3 the candidates from the e-:above3 list have the highest quality value
86
87# Merging the three lists with the e-:all list as first input list
88ma.copyLists('e-:onlyall', ['e-:allMC', 'e-:below2', 'e-:above3'], path=path)
89# Since e-:all is a superset of the other two lists, the merged list should contain only particles with identifier equals 0.
90ma.printVariableValues('e-:onlyall', ['extraInfo(identifier)'], path=path)
91# ma.applyCuts('e-:onlyall', 'extraInfo(identifier) == 0', path=path)
92
93# Merging the three lists with the e-:all list at the end
94ma.copyLists('e-:goodmerge', ['e-:below2', 'e-:above3', 'e-:allMC'], path=path)
95# All electrons from the e-:below2 and e-:above3 lists should be present in
96# the merged list, and those with 2 < px < 3 from the e-:all list.
97# Thus, the merged list contains the ten best (according to quality)
98# particles. This desired outcome was achieved accidentally.
99ma.printVariableValues('e-:goodmerge', ['extraInfo(identifier)'], path=path)
100ma.applyCuts('e-:goodmerge', 'extraInfo(identifier) > 0', path=path)
101
102# Merging the three lists with the e-:all list as first input list but providing the quality extra info as sorting variable
103path.add_module('ParticleListManipulator', outputListName='e-:best', inputListNames=['e-:allMC', 'e-:below2', 'e-:above3'],
104 variable='extraInfo(quality)', preferLowest=False)
105# The merged list should again contain the ten best particles. This time, independent of the order in the input list.
106ma.printVariableValues('e-:best', ['extraInfo(identifier)'], path=path)
107ma.applyCuts('e-:best', 'extraInfo(identifier) > 0', path=path)
108
109path.add_module(MergingChecker())
110
111basf2.process(path)
112
113# @endcond
A (simplified) python wrapper for StoreArray.
Definition: PyStoreArray.h:72
a (simplified) python wrapper for StoreObjPtr.
Definition: PyStoreObj.h:67