Belle II Software development
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 'ArmenterosLongitudinalMomentumAsymmetry',
65 'ArmenterosDaughter1Qt',
66 'ArmenterosDaughter2Qt'
67 ]
68 utils.add_collection(inputVariablesList, 'ks_selector_info')
69
70
71def V0Selector_Training(
72 train_data,
73 tree_name="tree",
74 mva_identifier="MVAFastBDT_V0Selector.root",
75 target_variable="isSignal",
76 parameters={}
77):
78 """
79 Defines the configuration of V0Selector Training.
80 The training data should contain K_S0 and misreconstructed K_S0 without Lambda0.
81
82 @param train_data Root file containing Ks information to be trained.
83 @param tree_name Tree name for variables.
84 @param mva_identifier Name for output MVA weight file.
85 @param target_variable Target variable for MVA training.
86 @param parameters hyperparameter for LGBM
87
88 """
89 trainVars = [
90 'cosVertexMomentum',
91 'flightDistance',
92 'significanceOfDistance',
93 'cosHelicityAngleMomentum',
94 'ImpactXY',
95 'decayAngle(0)',
96 'decayAngle(1)',
97 'daughterAngleDiffInMother',
98 'daughtersDeltaZ',
99 'pip_nSVDHits', 'pin_nSVDHits',
100 'pip_dr', 'pin_dr',
101 ]
102
103 general_options = basf2_mva.GeneralOptions()
104 general_options.m_datafiles = basf2_mva.vector(train_data)
105 general_options.m_treename = tree_name
106 general_options.m_identifier = mva_identifier
107 general_options.m_variables = basf2_mva.vector(*trainVars)
108 general_options.m_target_variable = target_variable
109 fastbdt_options = basf2_mva.FastBDTOptions()
110 basf2_mva.teacher(general_options, fastbdt_options)
111
112 general_options = basf2_mva.GeneralOptions()
113 general_options.m_datafiles = basf2_mva.vector(train_data)
114 general_options.m_treename = tree_name
115 general_options.m_identifier = mva_identifier+'.root'
116 general_options.m_variables = basf2_mva.vector(*trainVars)
117 general_options.m_target_variable = target_variable
118
119 python_options = basf2_mva.PythonOptions()
120
121 python_options.m_framework = "custom"
122 python_options.m_steering_file = "mva/scripts/basf2_mva_python_interface/lightgbm.py"
123
124 import json
125 param = {'num_leaves': 256,
126 'learning_rate': 0.1,
127 'device_type': "cpu",
128 'deterministic': True,
129 'metric': 'auc',
130 'num_round': 100,
131 # 'stop_round' : 30,
132 'path': mva_identifier+'.txt',
133 'max_bin': 250,
134 'boosting': ' gbdt',
135 'trainFraction': 0.8,
136 'min_data_in_leaf': 4000,
137 'max_depth': 8,
138 'objective': 'cross_entropy'
139 }
140
141 if isinstance(parameters, dict):
142 param.update(parameters)
143 config_string = json.dumps(param)
144 print("The json config string", config_string)
145 python_options.m_config = config_string
146
147 python_options.m_training_fraction = 1
148
149 python_options.m_normalize = False # we do it inside MVA torch
150
151 python_options.m_nIterations = 1
152 python_options.m_mini_batch_size = 0
153 basf2_mva.teacher(general_options, python_options)
154
155
156def LambdaVeto_Training(
157 train_data,
158 tree_name="tree",
159 mva_identifier="MVAFastBDT_LambdaVeto.root",
160 target_variable="isSignal",
161 parameters={}
162):
163 """
164 Defines the configuration of LambdaVeto Training.
165 The training data should contain only K_S0 and Lambda0.
166
167 @param train_data Root file containing Ks information to be trained.
168 @param tree_name Tree name for variables.
169 @param mva_identifier Name for output MVA weight file.
170 @param target_variable Target variable for MVA training.
171 @param parameters hyperparameter for LGBM
172 """
173 trainVars = [
174 'pip_protonID',
175 'pin_protonID',
176 'M_lambda_p',
177 'M_lambda_antip',
178 'pip_cosTheta',
179 'pin_cosTheta',
180 'ArmenterosLongitudinalMomentumAsymmetry',
181 'ArmenterosDaughter1Qt',
182 'ArmenterosDaughter2Qt'
183 ]
184 general_options = basf2_mva.GeneralOptions()
185 general_options.m_datafiles = basf2_mva.vector(train_data)
186 general_options.m_treename = tree_name
187 general_options.m_identifier = mva_identifier
188 general_options.m_variables = basf2_mva.vector(*trainVars)
189 general_options.m_target_variable = target_variable
190 fastbdt_options = basf2_mva.FastBDTOptions()
191 basf2_mva.teacher(general_options, fastbdt_options)
192
193 general_options = basf2_mva.GeneralOptions()
194 general_options.m_datafiles = basf2_mva.vector(train_data)
195 general_options.m_treename = tree_name
196 general_options.m_identifier = mva_identifier+'.root'
197 general_options.m_variables = basf2_mva.vector(*trainVars)
198 general_options.m_target_variable = target_variable
199
200 python_options = basf2_mva.PythonOptions()
201
202 python_options.m_framework = "custom"
203 python_options.m_steering_file = "mva/scripts/basf2_mva_python_interface/lightgbm.py"
204
205 import json
206 param = {'num_leaves': 256,
207 'learning_rate': 0.2,
208 'device_type': "cpu",
209 'deterministic': True,
210 'metric': 'auc',
211 'num_round': 100,
212 # 'stop_round' : 30,
213 'path': mva_identifier+'.txt',
214 'max_bin': 250,
215 'boosting': ' dart',
216 'trainFraction': 0.8,
217 'min_data_in_leaf': 300,
218 'max_depth': 8,
219 'objective': 'cross_entropy'
220 }
221
222 if isinstance(parameters, dict):
223 param.update(parameters)
224 config_string = json.dumps(param)
225 print("The json config string", config_string)
226 python_options.m_config = config_string
227
228 python_options.m_training_fraction = 1
229
230 python_options.m_normalize = False # we do it inside MVA torch
231
232 python_options.m_nIterations = 1
233 python_options.m_mini_batch_size = 0
234 basf2_mva.teacher(general_options, python_options)
235# ****************************************
236# KS Selector MAIN FUNCTION
237# ****************************************
238
239
240def ksSelector(
241 particleListName,
242 identifier_Ks="Ks_LGBM_V0Selector",
243 identifier_vLambda="Ks_LGBM_LambdaVeto",
244 output_label_name='',
245 extraInfoName_V0Selector='KsSelector_V0Selector',
246 extraInfoName_LambdaVeto='KsSelector_LambdaVeto',
247 useCustomThreshold=False,
248 threshold_V0Selector=0.90,
249 threshold_LambdaVeto=0.11,
250 path=None
251):
252 """
253 This function will apply K_S0 selection MVA on the given particleList.
254 By default this function appends MVA output as a extraInfo for the given particleList.
255 You can apply preset cut or custom cut by giving parameters. In this case,
256 a new particleList is created from the original particleList applying cuts on the MVA output.
257
258 @param particleLists Reconstructed Ks -> pi+ pi- list.
259 @param output_label_name Label of the returned Ks particleList.
260 When empty '', no cut is applied and new particleList is not created.
261 When custom name, the custom threshold is used, and useCustomThreshold
262 must be True.
263 When 'standard', 'tight', or 'loose', a cut with Ks efficiency
264 90%, 95%, and 85% is applied.
265 @param extraInfoName_V0Selector Variable name for V0Selector MVA output.
266 @param extraInfoName_LambdaVeto Variable name for LambdaVeto MVA output.
267 @param identifier_Ks Identifier name for V0Selector weight file.
268 @param identifier_vLambda Identifier name for LambdaVeto weight file.
269 @param useCustomThreshold Flag whether threshold_V0Selector and threshold_LambdaVeto are used.
270 @param threshold_V0Selector Threshold for V0Selector.
271 @param threshold_LambdaVeto Threshold for LambdaVeto.
272 @param path Basf2 path to execute.
273
274 """
275
276 add_default_ks_Selector_aliases()
277
278 path.add_module('MVAMultipleExperts',
279 listNames=[particleListName],
280 extraInfoNames=[extraInfoName_V0Selector, extraInfoName_LambdaVeto],
281 identifiers=[identifier_Ks, identifier_vLambda])
282
283 _effnames = ['standard', 'tight', 'loose']
284 outputListName = ''
285
286 if useCustomThreshold:
287 if output_label_name in _effnames:
288 B2ERROR('KsSelector: Specify label name except for \'standard\', \'tight\', and \'loose\' '
289 'when you use custom threshold.')
290 elif output_label_name == '':
291 B2ERROR('KsSelector: Specify label name when you use custom threshold.')
292 else:
293 outputListName = particleListName.split(':')[0] + ':' + output_label_name
294 B2INFO('KsSelector: Custom Cut is applied on '+outputListName+'.')
295 V0_thr = threshold_V0Selector
296 Lambda_thr = threshold_LambdaVeto
297 B2INFO('KsSelector: Threshold is (' + str(V0_thr) + ', ' + str(Lambda_thr) + ')')
298 cut_string = 'extraInfo('+extraInfoName_V0Selector+')>'+str(V0_thr) + \
299 ' and extraInfo('+extraInfoName_LambdaVeto+')>'+str(Lambda_thr)
300 ma.cutAndCopyLists(outputListName, particleListName, cut=cut_string, path=path)
301 else:
302 if output_label_name in _effnames:
303 outputListName = particleListName.split(':')[0] + ':' + output_label_name
304 V0_thr = 0
305 Lambda_thr = 0
306 if output_label_name == 'standard':
307 B2INFO('KsSelector: Standard Cut is applied on '+outputListName+'.')
308 V0_thr = 0.91
309 Lambda_thr = 0.19
310 elif output_label_name == 'tight':
311 B2INFO('KsSelector: Tight Cut is applied on '+outputListName+'.')
312 V0_thr = 0.97
313 Lambda_thr = 0.45
314 elif output_label_name == 'loose':
315 B2INFO('KsSelector: Loose Cut is applied on '+outputListName+'.')
316 V0_thr = 0.51
317 Lambda_thr = 0.02
318 B2INFO('KsSelector: Threshold is (' + str(V0_thr) + ', ' + str(Lambda_thr) + ')')
319 cut_string = 'extraInfo('+extraInfoName_V0Selector+')>'+str(V0_thr) + \
320 ' and extraInfo('+extraInfoName_LambdaVeto+')>'+str(Lambda_thr)
321 ma.cutAndCopyLists(outputListName, particleListName, cut=cut_string, path=path)
322 elif output_label_name == '':
323 outputListName = particleListName
324 else:
325 B2ERROR('KsSelector: Label should be \'\', \'standard\', \'tight\', or \'loose\' if you do'
326 'not apply custom threshold')
327
328 B2INFO('KsSelector: ParticleList '+outputListName+' is returned.')