15 import modularAnalysis
as ma
16 from variables
import variables
as vm
17 from ROOT
import TFile
21 """The unit test case for TreeFitter"""
24 """Run the test fit"""
26 testFile = tempfile.NamedTemporaryFile()
28 main = basf2.create_path()
31 ma.inputMdst(inputfile, path=main)
33 ma.fillParticleList(
'e+:a',
'electronID > 0.5', path=main)
34 ma.fillParticleList(
'gamma:a',
'0.05 < E < 1.0', path=main)
35 ma.correctBremsBelle(
'e+:corrected',
'e+:a',
'gamma:a', path=main)
37 ma.reconstructDecay(
'J/psi:corrected -> e+:corrected e-:corrected ?addbrems',
'', path=main)
38 ma.matchMCTruth(
'J/psi:corrected', path=main)
40 ma.applyCuts(
'J/psi:corrected',
41 'daughterSumOf(extraInfo(bremsCorrected)) == 2 and isSignal==1',
43 ma.applyEventCuts(
'nParticlesInList(J/psi:corrected)', path=main)
46 main.add_module(
'TreeFitter',
47 particleList=
'J/psi:corrected',
49 updateAllDaughters=
True)
51 ma.reconstructDecay(
'J/psi:no -> e+:a e-:a',
'', path=main)
52 ma.matchMCTruth(
'J/psi:no', path=main)
53 ma.applyCuts(
'J/psi:no',
'isSignal==1', path=main)
54 main.add_module(
'TreeFitter',
55 particleList=
'J/psi:no',
57 updateAllDaughters=
True)
59 ma.variablesToEventExtraInfo(
'J/psi:no', {
'chiProb':
'chiProb_woCorrection',
'dM':
'dM_woCorrection'}, path=main)
60 vm.addAlias(
'chiProb_woCorrection',
'eventExtraInfo(chiProb_woCorrection)')
61 vm.addAlias(
'dM_woCorrection',
'eventExtraInfo(dM_woCorrection)')
65 xyz_error = ['x_uncertainty', 'y_uncertainty', 'z_uncertainty']
67 vm.addAlias(f'e0_diff_{v}', f'formula(daughter(0, {v}) - daughter(0, daughter(0, {v})))')
68 vm.addAlias(f'e1_diff_{v}', f'formula(daughter(1, {v}) - daughter(1, daughter(0, {v})))')
69 vm.addAlias(f'e0_pull_{v}', f'formula(e0_diff_{v} / daughter(0, {v}))')
70 vm.addAlias(f'e1_pull_{v}', f'formula(e1_diff_{v} / daughter(1, {v}))')
71 alias_list += [f'e0_pull_{v}', f'e1_pull_{v}']
74 ntupler = basf2.register_module(
'VariablesToNtuple')
75 ntupler.param(
'fileName', testFile.name)
76 ntupler.param(
'variables',
77 [
'chiProb',
'chiProb_woCorrection',
'dM',
'dM_woCorrection'])
78 ntupler.param(
'particleList',
'J/psi:corrected')
79 main.add_module(ntupler)
81 basf2.process(main, 10000)
83 ntuplefile = TFile(testFile.name)
84 ntuple = ntuplefile.Get(
'ntuple')
86 self.assertFalse(ntuple.GetEntries() == 0,
"Ntuple is empty.")
88 dMImproved = ntuple.GetEntries(
"abs(dM) < abs(dM_woCorrection)")
90 chiProbImproved = ntuple.GetEntries(
"(chiProb > chiProb_woCorrection) && (abs(dM) < abs(dM_woCorrection))")
91 chiProbWorsened = ntuple.GetEntries(
"(chiProb < chiProb_woCorrection) && (abs(dM) < abs(dM_woCorrection))")
93 self.assertTrue(dMImproved == 229,
"Brems-correction algorithm or setting have been changed.")
95 print(f
"Vertex fit quality improved due to the brems-correction, # of candidates = {chiProbImproved}")
96 print(f
"Vertex fit quality worsened due to the brems-correction, # of candidates = {chiProbWorsened}")
98 self.assertTrue(chiProbImproved == dMImproved,
"Brems-correction causes negative effect on the vertex-fit.")
99 print(
"Test passed, cleaning up.")
102 if __name__ ==
'__main__':
def require_file(filename, data_type="", py_case=None)
def clean_working_directory()