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