Belle II Software development
treeFit_referenced_bremscorrection.py
1#!/usr/bin/env python3
2
3
10import unittest
11import tempfile
12import basf2
13import b2test_utils
14import modularAnalysis as ma
15from variables import variables as vm
16from ROOT import TFile
17
18
19class TestTreeFits(unittest.TestCase):
20 """The unit test case for TreeFitter"""
21
22 def testFit(self):
23 """Run the test fit"""
24
25 testFile = tempfile.NamedTemporaryFile()
26
27 main = basf2.create_path()
28
29 inputfile = b2test_utils.require_file('JPsi2ee_e2egamma.root', 'examples', py_case=self)
30 ma.inputMdst(inputfile, path=main)
31
32 ma.fillParticleList('e+:a', 'electronID > 0.5', path=main)
33 ma.fillParticleList('gamma:a', '0.05 < E < 1.0', path=main)
34 ma.correctBremsBelle('e+:corrected', 'e+:a', 'gamma:a', path=main)
35
36 ma.reconstructDecay('J/psi:corrected -> e+:corrected e-:corrected ?addbrems', '', path=main)
37 ma.matchMCTruth('J/psi:corrected', path=main)
38
39 ma.applyCuts('J/psi:corrected',
40 'daughterSumOf(extraInfo(bremsCorrected)) == 2 and isSignal==1',
41 path=main)
42 ma.applyEventCuts('nParticlesInList(J/psi:corrected)', path=main)
43
44 conf = -1
45 main.add_module('TreeFitter',
46 particleList='J/psi:corrected',
47 confidenceLevel=conf,
48 updateAllDaughters=True)
49
50 ma.reconstructDecay('J/psi:no -> e+:a e-:a', '', path=main)
51 ma.matchMCTruth('J/psi:no', path=main)
52 ma.applyCuts('J/psi:no', 'isSignal==1', path=main)
53 main.add_module('TreeFitter',
54 particleList='J/psi:no',
55 confidenceLevel=conf,
56 updateAllDaughters=True)
57
58 ma.variablesToEventExtraInfo('J/psi:no', {'chiProb': 'chiProb_woCorrection', 'dM': 'dM_woCorrection'}, path=main)
59 vm.addAlias('chiProb_woCorrection', 'eventExtraInfo(chiProb_woCorrection)')
60 vm.addAlias('dM_woCorrection', 'eventExtraInfo(dM_woCorrection)')
61
62 """
63 alias_list = []
64 xyz_error = ['x_uncertainty', 'y_uncertainty', 'z_uncertainty']
65 for v in xyz_error:
66 vm.addAlias(f'e0_diff_{v}', f'formula(daughter(0, {v}) - daughter(0, daughter(0, {v})))')
67 vm.addAlias(f'e1_diff_{v}', f'formula(daughter(1, {v}) - daughter(1, daughter(0, {v})))')
68 vm.addAlias(f'e0_pull_{v}', f'formula(e0_diff_{v} / daughter(0, {v}))')
69 vm.addAlias(f'e1_pull_{v}', f'formula(e1_diff_{v} / daughter(1, {v}))')
70 alias_list += [f'e0_pull_{v}', f'e1_pull_{v}']
71 """
72
73 ntupler = basf2.register_module('VariablesToNtuple')
74 ntupler.param('fileName', testFile.name)
75 ntupler.param('variables',
76 ['chiProb', 'chiProb_woCorrection', 'dM', 'dM_woCorrection'])
77 ntupler.param('particleList', 'J/psi:corrected')
78 main.add_module(ntupler)
79
80 basf2.process(main, 10000)
81
82 ntuplefile = TFile(testFile.name)
83 ntuple = ntuplefile.Get('ntuple')
84
85 self.assertFalse(ntuple.GetEntries() == 0, "Ntuple is empty.")
86
87 dMImproved = ntuple.GetEntries("abs(dM) < abs(dM_woCorrection)")
88
89 chiProbImproved = ntuple.GetEntries("(chiProb > chiProb_woCorrection) && (abs(dM) < abs(dM_woCorrection))")
90 chiProbWorsened = ntuple.GetEntries("(chiProb < chiProb_woCorrection) && (abs(dM) < abs(dM_woCorrection))")
91
92 self.assertTrue(dMImproved == 229, "Brems-correction algorithm or setting have been changed.")
93
94 print(f"Vertex fit quality improved due to the brems-correction, # of candidates = {chiProbImproved}")
95 print(f"Vertex fit quality worsened due to the brems-correction, # of candidates = {chiProbWorsened}")
96
97 self.assertTrue(chiProbImproved == dMImproved, "Brems-correction causes negative effect on the vertex-fit.")
98 print("Test passed, cleaning up.")
99
100
101if __name__ == '__main__':
103 unittest.main()
def require_file(filename, data_type="", py_case=None)
Definition: __init__.py:54
def clean_working_directory()
Definition: __init__.py:189