Belle II Software light-2406-ragdoll
test_particleExtractorFromROE.py
1#!/usr/bin/env python3
2
3
10
11import unittest
12import tempfile
13import basf2
14import b2test_utils
15import modularAnalysis as ma
16from ROOT import TFile, Belle2
17
18
19class TestParticleExtractorFromROE(unittest.TestCase):
20 """The unit test for ParticleExtractorFromROE module """
21
22 def test_roepath(self):
23 """ Test the extractParticlesFromROE function in roe_path """
24
25 testFile = tempfile.NamedTemporaryFile()
26
27 main = basf2.create_path()
28
29 inputfile = b2test_utils.require_file('analysis/B0ToPiPiPi0.root', 'validation', py_case=self)
30 ma.inputMdst(inputfile, path=main)
31
32 ma.fillParticleList('pi+:sig', 'pionID > 0.5', path=main)
33
34 ma.fillParticleList('gamma:all', '', path=main)
35 ma.reconstructDecay('pi0:sig -> gamma:all gamma:all', '0.125 < InvM < 0.145', path=main)
36
37 ma.reconstructDecay('B0:sig -> pi-:sig pi+:sig pi0:sig', '', path=main)
38
39 ma.matchMCTruth('B0:sig', path=main)
40 ma.applyCuts('B0:sig', 'isSignal==1', path=main)
41
42 ma.buildRestOfEvent('B0:sig', fillWithMostLikely=True, path=main)
43
44 roe_path = basf2.create_path()
45 deadEndPath = basf2.create_path()
46 ma.signalSideParticleFilter('B0:sig', '', roe_path, deadEndPath)
47
48 plists = [f'{ptype}:in_roe' for ptype in ['pi+', 'gamma', 'K_L0', 'K+', 'p+', 'e+', 'mu+']]
49 ma.extractParticlesFromROE(plists, maskName='all', path=roe_path)
50
51 charged_inROE = [f'nParticlesInList({ptype}:in_roe)' for ptype in ['pi+', 'K+', 'p+', 'e+', 'mu+']]
52 neutral_inROE = [f'nParticlesInList({ptype}:in_roe)' for ptype in ['gamma', 'K_L0']]
53
54 default = ['nParticlesInList(pi+:all)', 'nParticlesInList(gamma:all)', 'nParticlesInList(K_L0:roe_default)']
55 mostLikely = [f'nParticlesInList({ptype}:mostlikely_roe)' for ptype in ['K+', 'p+', 'e+', 'mu+']]
56
57 ma.variablesToNtuple('', charged_inROE + neutral_inROE + default + mostLikely,
58 filename=testFile.name,
59 path=roe_path)
60
61 main.for_each('RestOfEvent', 'RestOfEvents', roe_path)
62
63 basf2.process(main)
64
65 ntuplefile = TFile(testFile.name)
66 ntuple = ntuplefile.Get('variables')
67
68 self.assertFalse(ntuple.GetEntries() == 0, "Ntuple is empty.")
69
70 nEntries = ntuple.GetEntries()
71
72 # Number of pi+:all = Number of charged-FSPs in ROE + 2-pions in signal side
73 cut = Belle2.MakeROOTCompatible.makeROOTCompatible('nParticlesInList(pi+:all)') + ' == ' \
74 '(' + '+'.join([Belle2.MakeROOTCompatible.makeROOTCompatible(charged) for charged in charged_inROE]) + ' + 2)'
75 nPass_chargedCheck = ntuple.GetEntries(cut)
76 self.assertFalse(nPass_chargedCheck < nEntries, "Charged particles are not correctly extracted")
77
78 # Number of mostLikely = Number of charged in ROE
79 cut = Belle2.MakeROOTCompatible.makeROOTCompatible('nParticlesInList(e+:mostlikely_roe)') + ' == ' \
80 + Belle2.MakeROOTCompatible.makeROOTCompatible('nParticlesInList(e+:in_roe)')
81 nPass_electronCheck = ntuple.GetEntries(cut)
82 self.assertFalse(nPass_electronCheck < nEntries, "MostLikely option seems broken")
83
84 def test_mainpath(self):
85 """ Test the extractParticlesFromROE function in main path """
86
87 testFile = tempfile.NamedTemporaryFile()
88
89 main = basf2.create_path()
90
91 inputfile = b2test_utils.require_file('analysis/B0ToPiPiPi0.root', 'validation', py_case=self)
92 ma.inputMdst(inputfile, path=main)
93
94 ma.fillParticleList('pi+:sig', 'pionID > 0.5', path=main)
95
96 ma.fillParticleList('gamma:all', '', path=main)
97 ma.reconstructDecay('pi0:sig -> gamma:all gamma:all', '0.125 < InvM < 0.145', path=main)
98
99 ma.reconstructDecay('B0:sig -> pi-:sig pi+:sig pi0:sig', '', path=main)
100
101 ma.matchMCTruth('B0:sig', path=main)
102 ma.applyCuts('B0:sig', 'isSignal==1', path=main) # there must be only 1 candidate
103
104 ma.buildRestOfEvent('B0:sig', fillWithMostLikely=True, path=main)
105
106 plists = [f'{ptype}:in_roe' for ptype in ['pi+', 'gamma', 'K_L0', 'K+', 'p+', 'e+', 'mu+']]
107 ma.extractParticlesFromROE(plists, maskName='all', path=main, signalSideParticleList='B0:sig')
108
109 charged_inROE = [f'nParticlesInList({ptype}:in_roe)' for ptype in ['pi+', 'K+', 'p+', 'e+', 'mu+']]
110 neutral_inROE = [f'nParticlesInList({ptype}:in_roe)' for ptype in ['gamma', 'K_L0']]
111
112 default = ['nParticlesInList(pi+:all)', 'nParticlesInList(gamma:all)', 'nParticlesInList(K_L0:roe_default)']
113 mostLikely = [f'nParticlesInList({ptype}:mostlikely_roe)' for ptype in ['K+', 'p+', 'e+', 'mu+']]
114
115 ma.variablesToNtuple('B0:sig', charged_inROE + neutral_inROE + default + mostLikely,
116 filename=testFile.name,
117 path=main)
118
119 basf2.process(main)
120
121 ntuplefile = TFile(testFile.name)
122 ntuple = ntuplefile.Get('variables')
123
124 self.assertFalse(ntuple.GetEntries() == 0, "Ntuple is empty.")
125
126 nEntries = ntuple.GetEntries()
127
128 # Number of pi+:all = Number of charged-FSPs in ROE + 2-pions in signal side
129 cut = Belle2.MakeROOTCompatible.makeROOTCompatible('nParticlesInList(pi+:all)') + ' == ' \
130 '(' + '+'.join([Belle2.MakeROOTCompatible.makeROOTCompatible(charged) for charged in charged_inROE]) + ' + 2)'
131 nPass_chargedCheck = ntuple.GetEntries(cut)
132 self.assertFalse(nPass_chargedCheck < nEntries, "Charged particles are not correctly extracted")
133
134 # Number of mostLikely = Number of charged in ROE
135 cut = Belle2.MakeROOTCompatible.makeROOTCompatible('nParticlesInList(e+:mostlikely_roe)') + ' == ' \
136 + Belle2.MakeROOTCompatible.makeROOTCompatible('nParticlesInList(e+:in_roe)')
137 nPass_electronCheck = ntuple.GetEntries(cut)
138 self.assertFalse(nPass_electronCheck < nEntries, "MostLikely option seems broken")
139
140
141if __name__ == '__main__':
143 unittest.main()
static std::string makeROOTCompatible(std::string str)
Remove special characters that ROOT dislikes in branch names, e.g.
def require_file(filename, data_type="", py_case=None)
Definition: __init__.py:54
def clean_working_directory()
Definition: __init__.py:189