Belle II Software light-2406-ragdoll
test_DistanceCalculator.py
1
8import basf2 as b2
9import modularAnalysis as ma
10import variables.collections as vc
11import variables.utils as vu
12import vertex as vx
13import stdCharged as stdc
14import b2test_utils
15
16from variables import variables
17
18variables.addAlias('TubeX', 'extraInfo(TubePosX)')
19variables.addAlias('TubeY', 'extraInfo(TubePosY)')
20variables.addAlias('TubeZ', 'extraInfo(TubePosZ)')
21variables.addAlias('Tube00', 'extraInfo(TubeCov00)')
22variables.addAlias('Tube01', 'extraInfo(TubeCov01)')
23variables.addAlias('Tube02', 'extraInfo(TubeCov02)')
24variables.addAlias('Tube10', 'extraInfo(TubeCov10)')
25variables.addAlias('Tube11', 'extraInfo(TubeCov11)')
26variables.addAlias('Tube12', 'extraInfo(TubeCov12)')
27variables.addAlias('Tube20', 'extraInfo(TubeCov20)')
28variables.addAlias('Tube21', 'extraInfo(TubeCov21)')
29variables.addAlias('Tube22', 'extraInfo(TubeCov22)')
30
31variables.addAlias('avf00', 'extraInfo(prod_vtx_cov00)')
32variables.addAlias('avf01', 'extraInfo(prod_vtx_cov01)')
33variables.addAlias('avf02', 'extraInfo(prod_vtx_cov02)')
34variables.addAlias('avf10', 'extraInfo(prod_vtx_cov10)')
35variables.addAlias('avf11', 'extraInfo(prod_vtx_cov11)')
36variables.addAlias('avf12', 'extraInfo(prod_vtx_cov12)')
37variables.addAlias('avf20', 'extraInfo(prod_vtx_cov20)')
38variables.addAlias('avf21', 'extraInfo(prod_vtx_cov21)')
39variables.addAlias('avf22', 'extraInfo(prod_vtx_cov22)')
40
41variables.addAlias('avfPx', 'extraInfo(Px_after_avf)')
42variables.addAlias('avfPy', 'extraInfo(Py_after_avf)')
43variables.addAlias('avfPz', 'extraInfo(Pz_after_avf)')
44variables.addAlias('avfE', 'extraInfo(E_after_avf)')
45
46variables.addAlias('mcPxCms', 'useCMSFrame(mcPX)')
47variables.addAlias('mcPyCms', 'useCMSFrame(mcPY)')
48variables.addAlias('mcPzCms', 'useCMSFrame(mcPZ)')
49variables.addAlias('mcECms', 'useCMSFrame(mcE)')
50
51variables.addAlias('tube_dir_X', 'extraInfo(TubeDirX)')
52variables.addAlias('tube_dir_Y', 'extraInfo(TubeDirY)')
53variables.addAlias('tube_dir_Z', 'extraInfo(TubeDirZ)')
54
55variables.addAlias('truex', 'matchedMC(x)')
56variables.addAlias('truey', 'matchedMC(y)')
57variables.addAlias('truez', 'matchedMC(z)')
58
59variables.addAlias('Distance', 'extraInfo(CalculatedDistance)')
60variables.addAlias('DistanceError', 'extraInfo(CalculatedDistanceError)')
61variables.addAlias('DistanceVector_X', 'extraInfo(CalculatedDistanceVector_X)')
62variables.addAlias('DistanceVector_Y', 'extraInfo(CalculatedDistanceVector_Y)')
63variables.addAlias('DistanceVector_Z', 'extraInfo(CalculatedDistanceVector_Z)')
64variables.addAlias('DistanceCovMatrixXX', 'extraInfo(CalculatedDistanceCovMatrixXX)')
65variables.addAlias('DistanceCovMatrixXY', 'extraInfo(CalculatedDistanceCovMatrixXY)')
66variables.addAlias('DistanceCovMatrixXZ', 'extraInfo(CalculatedDistanceCovMatrixXZ)')
67variables.addAlias('DistanceCovMatrixYX', 'extraInfo(CalculatedDistanceCovMatrixYX)')
68variables.addAlias('DistanceCovMatrixYY', 'extraInfo(CalculatedDistanceCovMatrixYY)')
69variables.addAlias('DistanceCovMatrixYZ', 'extraInfo(CalculatedDistanceCovMatrixYZ)')
70variables.addAlias('DistanceCovMatrixZX', 'extraInfo(CalculatedDistanceCovMatrixZX)')
71variables.addAlias('DistanceCovMatrixZY', 'extraInfo(CalculatedDistanceCovMatrixZY)')
72variables.addAlias('DistanceCovMatrixZZ', 'extraInfo(CalculatedDistanceCovMatrixZZ)')
73# create path
74my_path = b2.create_path()
75
76# load input ROOT file
77
78ma.inputMdst(b2test_utils.require_file('analysis/tests/mdst.root'),
79 path=my_path)
80
81stdc.stdPi(listtype='all', path=my_path)
82stdc.stdPi(listtype='good', path=my_path)
83stdc.stdK(listtype='all', path=my_path)
84stdc.stdK(listtype='good', path=my_path)
85stdc.stdMu(listtype='all', path=my_path)
86stdc.stdMu(listtype='good', path=my_path)
87
88ma.reconstructDecay('D0:kpi -> K-:good pi+:good', '', path=my_path)
89ma.matchMCTruth('D0:kpi', path=my_path)
90ma.applyCuts('D0:kpi', 'abs(mcPDG)==421', path=my_path)
91vx.raveFit('D0:kpi', 0.0, path=my_path)
92
93ma.reconstructDecay('J/psi:mumu -> mu+:good mu-:good', '', path=my_path)
94vx.raveFit('J/psi:mumu', 0.0, path=my_path)
95
96# TAGGING
97ma.reconstructDecay('B+:tag -> anti-D0:kpi pi+:all', '', path=my_path)
98ma.matchMCTruth('B+:tag', path=my_path)
99ma.applyCuts('B+:tag', 'isSignalAcceptMissingNeutrino>0.5', path=my_path)
100vx.raveFit('B+:tag', 0.0, path=my_path)
101
102# SIGNAL
103ma.reconstructDecay('B-:sigT -> J/psi:mumu K-:good pi+:good pi-:good', '', path=my_path)
104ma.matchMCTruth('B-:sigT', path=my_path)
105ma.applyCuts('B-:sigT', 'abs(mcPDG)==521', path=my_path)
106
107ma.reconstructDecay('Upsilon(4S):sig -> B+:tag B-:sigT', '', path=my_path)
108
109mytestmodule2 = b2.register_module('BtubeCreator')
110my_path.add_module(mytestmodule2,
111 listName='Upsilon(4S):sig',
112 decayString='Upsilon(4S) -> ^B+ B-')
113vx.raveFit('B-:sigT', 0.0, constraint='btube', path=my_path)
114
115ma.calculateDistance('B-:sigT', 'B- -> ^J/psi K- pi+ pi-', "vertexbtube", path=my_path)
116# Saving variables to ntuple
117output_file = 'test.root'
118true_vars = ['truex', 'truey', 'truez']
119avf_vars = ['avf00', 'avf01', 'avf02', 'avf10', 'avf11', 'avf12', 'avf20', 'avf21', 'avf22', 'avfPx', 'avfPy', 'avfPz', 'avfE']
120tube_vars = [
121 'TubeX',
122 'TubeY',
123 'TubeZ',
124 'Tube00',
125 'Tube01',
126 'Tube02',
127 'Tube10',
128 'Tube11',
129 'Tube12',
130 'Tube20',
131 'Tube21',
132 'Tube22',
133 'tube_dir_X',
134 'tube_dir_Y',
135 'tube_dir_Z']
136
137common_vars = vc.kinematics + vc.vertex + vc.mc_vertex + vc.mc_truth + \
138 vc.mc_kinematics + ['InvM', 'mcPxCms', 'mcPyCms', 'mcPzCms', 'mcECms']
139distance_vars = [
140 'Distance',
141 'DistanceError',
142 'DistanceVector_X',
143 'DistanceVector_Y',
144 'DistanceVector_Z',
145 'DistanceCovMatrixXX',
146 'DistanceCovMatrixXY',
147 'DistanceCovMatrixXZ',
148 'DistanceCovMatrixYX',
149 'DistanceCovMatrixYY',
150 'DistanceCovMatrixYZ',
151 'DistanceCovMatrixZX',
152 'DistanceCovMatrixZY',
153 'DistanceCovMatrixZZ']
154Distance_vars_sig = vu.create_aliases_for_selected(
155 list_of_variables=distance_vars,
156 decay_string='Upsilon(4S):sig -> B+:tag ^B-:sigT',
157 prefix='Distance')
158D_vars_tagside = vu.create_aliases_for_selected(
159 list_of_variables=common_vars,
160 decay_string='B+:tag -> ^anti-D0:kpi pi+:all',
161 prefix='D_tagside')
162U4S_vars_sig = vu.create_aliases_for_selected(
163 list_of_variables=common_vars + tube_vars + true_vars,
164 decay_string='Upsilon(4S):sig -> B+:tag ^B-:sigT',
165 prefix='sigB')
166U4S_vars_tag = vu.create_aliases_for_selected(
167 list_of_variables=common_vars + avf_vars + true_vars,
168 decay_string='Upsilon(4S):sig -> ^B+:tag B-:sigT',
169 prefix='tagB')
170U4S_vars = common_vars + U4S_vars_tag + U4S_vars_sig + D_vars_tagside + Distance_vars_sig
171ma.variablesToNtuple('Upsilon(4S):sig', U4S_vars,
172 filename=output_file, treename='u4stree', path=my_path)
173ma.variablesToNtuple('B+:tag', common_vars,
174 filename=output_file, treename='tagtree', path=my_path)
175with b2test_utils.clean_working_directory():
176 assert b2test_utils.safe_process(my_path) == 0
b2test_utils.require_file
def require_file(filename, data_type="", py_case=None)
Definition: __init__.py:54
b2test_utils.clean_working_directory
def clean_working_directory()
Definition: __init__.py:189
b2test_utils.safe_process
def safe_process(*args, **kwargs)
Definition: __init__.py:236