Belle II Software light-2406-ragdoll
ksSelector.py
1#!/usr/bin/env python3
2
3
10
11from basf2 import B2INFO, B2ERROR
12import basf2_mva
13import modularAnalysis as ma
14import variables
15from variables import utils
16
17
18def add_default_ks_Selector_aliases():
19 """
20 This function is used to set up variables aliases for ks Selector variables.
21 """
22 B2INFO('KsSelector: creating variables alias.')
23 variables.variables.addAlias('M_lambda_p', 'useAlternativeDaughterHypothesis(M, 0:p+)')
24 variables.variables.addAlias('M_lambda_antip', 'useAlternativeDaughterHypothesis(M, 1:anti-p-)')
25 variables.variables.addAlias('daughtersDeltaZ', 'daughterDiffOf(0, 1, dz)')
26 variables.variables.addAlias('cosVertexMomentum', 'cosAngleBetweenMomentumAndVertexVector')
27 variables.variables.addAlias('pip_nPXDHits', 'daughter(0,nPXDHits)')
28 variables.variables.addAlias('pin_nPXDHits', 'daughter(1,nPXDHits)')
29 variables.variables.addAlias('pip_nSVDHits', 'daughter(0,nSVDHits)')
30 variables.variables.addAlias('pin_nSVDHits', 'daughter(1,nSVDHits)')
31 variables.variables.addAlias('daughterAngleDiffInMother', 'useRestFrame(daughterAngle(0, 1))')
32 variables.variables.addAlias('pip_p', 'daughter(0,p)')
33 variables.variables.addAlias('pin_p', 'daughter(1,p)')
34 variables.variables.addAlias('pip_dr', 'daughter(0,dr)')
35 variables.variables.addAlias('pin_dr', 'daughter(1,dr)')
36 variables.variables.addAlias('pip_cosTheta', 'daughter(0,cosTheta)')
37 variables.variables.addAlias('pin_cosTheta', 'daughter(1,cosTheta)')
38 variables.variables.addAlias('pip_protonID', 'daughter(0,protonID)')
39 variables.variables.addAlias('pin_protonID', 'daughter(1,protonID)')
40
41
42def add_variable_collection():
43 """
44 Call this function to add variable collection for ksSelector.
45 """
46 add_default_ks_Selector_aliases()
47 inputVariablesList = [
48 'cosVertexMomentum',
49 'flightDistance',
50 'significanceOfDistance',
51 'cosHelicityAngleMomentum',
52 'ImpactXY',
53 'decayAngle(0)',
54 'decayAngle(1)',
55 'daughterAngleDiffInMother',
56 'daughtersDeltaZ',
57 'pip_nSVDHits', 'pip_nPXDHits',
58 'pin_nSVDHits', 'pin_nPXDHits',
59 'pip_dr', 'pin_dr',
60 'pip_protonID', 'pin_protonID',
61 'M_lambda_p', 'M_lambda_antip',
62 'pip_p', 'pin_p',
63 'pip_cosTheta', 'pin_cosTheta',
64 ]
65 utils.add_collection(inputVariablesList, 'ks_selector_info')
66
67
68def V0Selector_Training(
69 train_data,
70 tree_name="tree",
71 mva_identifier="MVAFastBDT_V0Selector.root",
72 target_variable="isSignal"
73):
74 """
75 Defines the configuration of V0Selector Training.
76 The training data should contain K_S0 and misreconstructed K_S0 without Lambda0.
77
78 @param train_data Root file containing Ks information to be trained.
79 @param tree_name Tree name for variables.
80 @param mva_identifier Name for output MVA weight file.
81 @param target_variable Target variable for MVA training.
82 """
83 trainVars = [
84 'cosVertexMomentum',
85 'flightDistance',
86 'significanceOfDistance',
87 'cosHelicityAngleMomentum',
88 'ImpactXY',
89 'decayAngle(0)',
90 'decayAngle(1)',
91 'daughterAngleDiffInMother',
92 'daughtersDeltaZ',
93 'pip_nSVDHits', 'pip_nPXDHits',
94 'pin_nSVDHits', 'pin_nPXDHits',
95 'pip_dr', 'pin_dr',
96 ]
97
98 general_options = basf2_mva.GeneralOptions()
99 general_options.m_datafiles = basf2_mva.vector(train_data)
100 general_options.m_treename = tree_name
101 general_options.m_identifier = mva_identifier
102 general_options.m_variables = basf2_mva.vector(*trainVars)
103 general_options.m_target_variable = target_variable
104 fastbdt_options = basf2_mva.FastBDTOptions()
105 basf2_mva.teacher(general_options, fastbdt_options)
106
107
108def LambdaVeto_Training(
109 train_data,
110 tree_name="tree",
111 mva_identifier="MVAFastBDT_LambdaVeto.root",
112 target_variable="isSignal"
113):
114 """
115 Defines the configuration of LambdaVeto Training.
116 The training data should contain only K_S0 and Lambda0.
117
118 @param train_data Root file containing Ks information to be trained.
119 @param tree_name Tree name for variables.
120 @param mva_identifier Name for output MVA weight file.
121 @param target_variable Target variable for MVA training.
122 """
123 trainVars = [
124 'pip_protonID',
125 'pin_protonID',
126 'M_lambda_p',
127 'M_lambda_antip',
128 'pip_p',
129 'pin_p',
130 'pip_cosTheta',
131 'pin_cosTheta',
132 ]
133 general_options = basf2_mva.GeneralOptions()
134 general_options.m_datafiles = basf2_mva.vector(train_data)
135 general_options.m_treename = tree_name
136 general_options.m_identifier = mva_identifier
137 general_options.m_variables = basf2_mva.vector(*trainVars)
138 general_options.m_target_variable = target_variable
139 fastbdt_options = basf2_mva.FastBDTOptions()
140 basf2_mva.teacher(general_options, fastbdt_options)
141
142# ****************************************
143# KS Selector MAIN FUNCTION
144# ****************************************
145
146
147def ksSelector(
148 particleListName,
149 identifier_Ks,
150 identifier_vLambda,
151 output_label_name='',
152 extraInfoName_V0Selector='KsSelector_V0Selector',
153 extraInfoName_LambdaVeto='KsSelector_LambdaVeto',
154 useCustomThreshold=False,
155 threshold_V0Selector=0.90,
156 threshold_LambdaVeto=0.11,
157 path=None
158):
159 """
160 This function will apply K_S0 selection MVA on the given particleList.
161 By default this function appends MVA output as a extraInfo for the given particleList.
162 You can apply preset cut or custom cut by giving parameters. In this case,
163 a new particleList is created from the original particleList applying cuts on the MVA output.
164
165 @param particleLists Reconstructed Ks -> pi+ pi- list.
166 @param output_label_name Label of the returned Ks particleList.
167 When empty '', no cut is applied and new particleList is not created.
168 When custom name, the custom threshold is used, and useCustomThreshold
169 must be True.
170 When 'standard', 'tight', or 'loose', a cut with Ks efficiency
171 90%, 95%, and 85% is applied.
172 @param extraInfoName_V0Selector Variable name for V0Selector MVA output.
173 @param extraInfoName_LambdaVeto Variable name for LambdaVeto MVA output.
174 @param identifier_Ks Identifier name for V0Selector weight file.
175 @param identifier_vLambda Identifier name for LambdaVeto weight file.
176 @param useCustomThreshold Flag whether threshold_V0Selector and threshold_LambdaVeto are used.
177 @param threshold_V0Selector Threshold for V0Selector.
178 @param threshold_LambdaVeto Threshold for LambdaVeto.
179 @param path Basf2 path to execute.
180
181 """
182
183 add_default_ks_Selector_aliases()
184
185 path.add_module('MVAMultipleExperts',
186 listNames=[particleListName],
187 extraInfoNames=[extraInfoName_V0Selector, extraInfoName_LambdaVeto],
188 identifiers=[identifier_Ks, identifier_vLambda])
189
190 _effnames = ['standard', 'tight', 'loose']
191 outputListName = ''
192
193 if useCustomThreshold:
194 if output_label_name in _effnames:
195 B2ERROR('KsSelector: Specify label name except for \'standard\', \'tight\', and \'loose\' '
196 'when you use custom threshold.')
197 elif output_label_name == '':
198 B2ERROR('KsSelector: Specify label name when you use custom threshold.')
199 else:
200 outputListName = particleListName.split(':')[0] + ':' + output_label_name
201 B2INFO('KsSelector: Custom Cut is applied on '+outputListName+'.')
202 V0_thr = threshold_V0Selector
203 Lambda_thr = threshold_LambdaVeto
204 B2INFO('KsSelector: Threshold is (' + str(V0_thr) + ', ' + str(Lambda_thr) + ')')
205 cut_string = 'extraInfo('+extraInfoName_V0Selector+')>'+str(V0_thr) + \
206 ' and extraInfo('+extraInfoName_LambdaVeto+')>'+str(Lambda_thr)
207 ma.cutAndCopyLists(outputListName, particleListName, cut=cut_string, path=path)
208 else:
209 if output_label_name in _effnames:
210 outputListName = particleListName.split(':')[0] + ':' + output_label_name
211 V0_thr = 0
212 Lambda_thr = 0
213 if output_label_name == 'standard':
214 B2INFO('KsSelector: Standard Cut is applied on '+outputListName+'.')
215 V0_thr = 0.90
216 Lambda_thr = 0.11
217 elif output_label_name == 'tight':
218 B2INFO('KsSelector: Tight Cut is applied on '+outputListName+'.')
219 V0_thr = 0.96
220 Lambda_thr = 0.27
221 elif output_label_name == 'loose':
222 B2INFO('KsSelector: Loose Cut is applied on '+outputListName+'.')
223 V0_thr = 0.49
224 Lambda_thr = 0.02
225 B2INFO('KsSelector: Threshold is (' + str(V0_thr) + ', ' + str(Lambda_thr) + ')')
226 cut_string = 'extraInfo('+extraInfoName_V0Selector+')>'+str(V0_thr) + \
227 ' and extraInfo('+extraInfoName_LambdaVeto+')>'+str(Lambda_thr)
228 ma.cutAndCopyLists(outputListName, particleListName, cut=cut_string, path=path)
229 elif output_label_name == '':
230 outputListName = particleListName
231 else:
232 B2ERROR('KsSelector: Label should be \'\', \'standard\', \'tight\', or \'loose\' if you do'
233 'not apply custom threshold')
234
235 B2INFO('KsSelector: ParticleList '+outputListName+' is returned.')