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