Belle II Software  light-2212-foldex
flavorTagger.py
1 #!/usr/bin/env python3
2 
3 
10 
11 # ************* Flavor Tagging ************
12 # * This script is needed to train *
13 # * and to test the flavor tagger. *
14 # ********************************************
15 
16 from basf2 import B2INFO, B2FATAL, B2WARNING
17 import basf2
18 import basf2_mva
19 import inspect
20 import modularAnalysis as ma
21 import variables
22 from variables import utils
23 import os
24 import glob
25 import b2bii
26 import collections
27 
28 
29 def getBelleOrBelle2():
30  """
31  Gets the global ModeCode.
32  """
33  if b2bii.isB2BII():
34  return 'Belle'
35  else:
36  return 'Belle2'
37 
38 
39 def setInteractionWithDatabase(downloadFromDatabaseIfNotFound=False, uploadToDatabaseAfterTraining=False):
40  """
41  Sets the interaction with the database: download trained weight files or upload weight files after training.
42  """
43 
44  global downloadFlag
45  global uploadFlag
46 
47  downloadFlag = downloadFromDatabaseIfNotFound
48  uploadFlag = uploadToDatabaseAfterTraining
49 
50 
51 # Default list of aliases that should be used to save the flavor tagging information using VariablesToNtuple
52 flavor_tagging = ['FBDT_qrCombined', 'FANN_qrCombined', 'qrMC', 'mcFlavorOfOtherB',
53  'qpElectron', 'hasTrueTargetElectron', 'isRightCategoryElectron',
54  'qpIntermediateElectron', 'hasTrueTargetIntermediateElectron', 'isRightCategoryIntermediateElectron',
55  'qpMuon', 'hasTrueTargetMuon', 'isRightCategoryMuon',
56  'qpIntermediateMuon', 'hasTrueTargetIntermediateMuon', 'isRightCategoryIntermediateMuon',
57  'qpKinLepton', 'hasTrueTargetKinLepton', 'isRightCategoryKinLepton',
58  'qpIntermediateKinLepton', 'hasTrueTargetIntermediateKinLepton', 'isRightCategoryIntermediateKinLepton',
59  'qpKaon', 'hasTrueTargetKaon', 'isRightCategoryKaon',
60  'qpSlowPion', 'hasTrueTargetSlowPion', 'isRightCategorySlowPion',
61  'qpFastHadron', 'hasTrueTargetFastHadron', 'isRightCategoryFastHadron',
62  'qpLambda', 'hasTrueTargetLambda', 'isRightCategoryLambda',
63  'qpFSC', 'hasTrueTargetFSC', 'isRightCategoryFSC',
64  'qpMaximumPstar', 'hasTrueTargetMaximumPstar', 'isRightCategoryMaximumPstar',
65  'qpKaonPion', 'hasTrueTargetKaonPion', 'isRightCategoryKaonPion']
66 
67 
68 def add_default_FlavorTagger_aliases():
69  """
70  This function adds the default aliases for flavor tagging variables
71  and defines the collection of flavor tagging variables.
72  """
73 
74  variables.variables.addAlias('FBDT_qrCombined', 'qrOutput(FBDT)')
75  variables.variables.addAlias('FANN_qrCombined', 'qrOutput(FANN)')
76  variables.variables.addAlias('qrMC', 'isRelatedRestOfEventB0Flavor')
77 
78  for iCategory in AvailableCategories:
79  aliasForQp = 'qp' + iCategory
80  aliasForTrueTarget = 'hasTrueTarget' + iCategory
81  aliasForIsRightCategory = 'isRightCategory' + iCategory
82  variables.variables.addAlias(aliasForQp, 'qpCategory(' + iCategory + ')')
83  variables.variables.addAlias(aliasForTrueTarget, 'hasTrueTargets(' + iCategory + ')')
84  variables.variables.addAlias(aliasForIsRightCategory, 'isTrueFTCategory(' + iCategory + ')')
85 
86  utils.add_collection(flavor_tagging, 'flavor_tagging')
87 
88 
89 def set_FlavorTagger_pid_aliases():
90  """
91  This function adds the pid aliases needed by the flavor tagger.
92  """
93  variables.variables.addAlias('eid_TOP', 'pidPairProbabilityExpert(11, 211, TOP)')
94  variables.variables.addAlias('eid_ARICH', 'pidPairProbabilityExpert(11, 211, ARICH)')
95  variables.variables.addAlias('eid_ECL', 'pidPairProbabilityExpert(11, 211, ECL)')
96 
97  variables.variables.addAlias('muid_TOP', 'pidPairProbabilityExpert(13, 211, TOP)')
98  variables.variables.addAlias('muid_ARICH', 'pidPairProbabilityExpert(13, 211, ARICH)')
99  variables.variables.addAlias('muid_KLM', 'pidPairProbabilityExpert(13, 211, KLM)')
100 
101  variables.variables.addAlias('piid_TOP', 'pidPairProbabilityExpert(211, 321, TOP)')
102  variables.variables.addAlias('piid_ARICH', 'pidPairProbabilityExpert(211, 321, ARICH)')
103 
104  variables.variables.addAlias('Kid_TOP', 'pidPairProbabilityExpert(321, 211, TOP)')
105  variables.variables.addAlias('Kid_ARICH', 'pidPairProbabilityExpert(321, 211, ARICH)')
106 
107  if getBelleOrBelle2() == "Belle":
108  variables.variables.addAlias('eid_dEdx', 'ifNANgiveX(pidPairProbabilityExpert(11, 211, CDC, SVD), 0.5)')
109  variables.variables.addAlias('muid_dEdx', 'ifNANgiveX(pidPairProbabilityExpert(13, 211, CDC, SVD), 0.5)')
110  variables.variables.addAlias('piid_dEdx', 'ifNANgiveX(pidPairProbabilityExpert(211, 321, CDC, SVD), 0.5)')
111  variables.variables.addAlias('pi_vs_edEdxid', 'ifNANgiveX(pidPairProbabilityExpert(211, 11, CDC, SVD), 0.5)')
112  variables.variables.addAlias('Kid_dEdx', 'ifNANgiveX(pidPairProbabilityExpert(321, 211, CDC, SVD), 0.5)')
113  else:
114  variables.variables.addAlias('eid_dEdx', 'pidPairProbabilityExpert(11, 211, CDC)')
115  variables.variables.addAlias('muid_dEdx', 'pidPairProbabilityExpert(13, 211, CDC)')
116  variables.variables.addAlias('piid_dEdx', 'pidPairProbabilityExpert(211, 321, CDC)')
117  variables.variables.addAlias('pi_vs_edEdxid', 'pidPairProbabilityExpert(211, 11, CDC)')
118  variables.variables.addAlias('Kid_dEdx', 'pidPairProbabilityExpert(321, 211, CDC)')
119 
120 
121 def set_FlavorTagger_pid_aliases_legacy():
122  """
123  This function adds the pid aliases needed by the flavor tagger trained for MC13.
124  """
125  variables.variables.addAlias('eid_TOP', 'ifNANgiveX(pidPairProbabilityExpert(11, 211, TOP), 0.5)')
126  variables.variables.addAlias('eid_ARICH', 'ifNANgiveX(pidPairProbabilityExpert(11, 211, ARICH), 0.5)')
127  variables.variables.addAlias('eid_ECL', 'ifNANgiveX(pidPairProbabilityExpert(11, 211, ECL), 0.5)')
128 
129  variables.variables.addAlias('muid_TOP', 'ifNANgiveX(pidPairProbabilityExpert(13, 211, TOP), 0.5)')
130  variables.variables.addAlias('muid_ARICH', 'ifNANgiveX(pidPairProbabilityExpert(13, 211, ARICH), 0.5)')
131  variables.variables.addAlias('muid_KLM', 'ifNANgiveX(pidPairProbabilityExpert(13, 211, KLM), 0.5)')
132 
133  variables.variables.addAlias('piid_TOP', 'ifNANgiveX(pidPairProbabilityExpert(211, 321, TOP), 0.5)')
134  variables.variables.addAlias('piid_ARICH', 'ifNANgiveX(pidPairProbabilityExpert(211, 321, ARICH), 0.5)')
135 
136  variables.variables.addAlias('Kid_TOP', 'ifNANgiveX(pidPairProbabilityExpert(321, 211, TOP), 0.5)')
137  variables.variables.addAlias('Kid_ARICH', 'ifNANgiveX(pidPairProbabilityExpert(321, 211, ARICH), 0.5)')
138 
139  if getBelleOrBelle2() == "Belle":
140  variables.variables.addAlias('eid_dEdx', 'ifNANgiveX(pidPairProbabilityExpert(11, 211, CDC, SVD), 0.5)')
141  variables.variables.addAlias('muid_dEdx', 'ifNANgiveX(pidPairProbabilityExpert(13, 211, CDC, SVD), 0.5)')
142  variables.variables.addAlias('piid_dEdx', 'ifNANgiveX(pidPairProbabilityExpert(211, 321, CDC, SVD), 0.5)')
143  variables.variables.addAlias('pi_vs_edEdxid', 'ifNANgiveX(pidPairProbabilityExpert(211, 11, CDC, SVD), 0.5)')
144  variables.variables.addAlias('Kid_dEdx', 'ifNANgiveX(pidPairProbabilityExpert(321, 211, CDC, SVD), 0.5)')
145  else:
146  # Removed SVD PID for Belle II MC and data as it is absent in release 4.
147  variables.variables.addAlias('eid_dEdx', 'ifNANgiveX(pidPairProbabilityExpert(11, 211, CDC), 0.5)')
148  variables.variables.addAlias('muid_dEdx', 'ifNANgiveX(pidPairProbabilityExpert(13, 211, CDC), 0.5)')
149  variables.variables.addAlias('piid_dEdx', 'ifNANgiveX(pidPairProbabilityExpert(211, 321, CDC), 0.5)')
150  variables.variables.addAlias('pi_vs_edEdxid', 'ifNANgiveX(pidPairProbabilityExpert(211, 11, CDC), 0.5)')
151  variables.variables.addAlias('Kid_dEdx', 'ifNANgiveX(pidPairProbabilityExpert(321, 211, CDC), 0.5)')
152 
153 
154 def setInputVariablesWithMask(maskName='all'):
155  """
156  Set aliases for input variables with ROE mask.
157  """
158  variables.variables.addAlias('pMissTag_withMask', 'pMissTag('+maskName+')')
159  variables.variables.addAlias('cosTPTO_withMask', 'cosTPTO('+maskName+')')
160  variables.variables.addAlias('ptTracksRoe_withMask', 'ptTracksRoe('+maskName+')')
161  variables.variables.addAlias('pt2TracksRoe_withMask', 'pt2TracksRoe('+maskName+')')
162  variables.variables.addAlias('ptTracksRoe_withMask', 'ptTracksRoe('+maskName+')')
163 
164 
165 def getFastBDTCategories():
166  '''
167  Helper function for getting the FastBDT categories.
168  It's necessary for removing top-level ROOT imports.
169  '''
170  fastBDTCategories = basf2_mva.FastBDTOptions()
171  fastBDTCategories.m_nTrees = 500
172  fastBDTCategories.m_nCuts = 8
173  fastBDTCategories.m_nLevels = 3
174  fastBDTCategories.m_shrinkage = 0.10
175  fastBDTCategories.m_randRatio = 0.5
176  return fastBDTCategories
177 
178 
179 def getFastBDTCombiner():
180  '''
181  Helper function for getting the FastBDT combiner.
182  It's necessary for removing top-level ROOT imports.
183  '''
184  fastBDTCombiner = basf2_mva.FastBDTOptions()
185  fastBDTCombiner.m_nTrees = 500
186  fastBDTCombiner.m_nCuts = 8
187  fastBDTCombiner.m_nLevels = 3
188  fastBDTCombiner.m_shrinkage = 0.10
189  fastBDTCombiner.m_randRatio = 0.5
190  return fastBDTCombiner
191 
192 
193 def getMlpFANNCombiner():
194  '''
195  Helper function for getting the MLP FANN combiner.
196  It's necessary for removing top-level ROOT imports.
197  '''
198  mlpFANNCombiner = basf2_mva.FANNOptions()
199  mlpFANNCombiner.m_max_epochs = 10000
200  mlpFANNCombiner.m_hidden_layers_architecture = "3*N"
201  mlpFANNCombiner.m_hidden_activiation_function = "FANN_SIGMOID_SYMMETRIC"
202  mlpFANNCombiner.m_output_activiation_function = "FANN_SIGMOID_SYMMETRIC"
203  mlpFANNCombiner.m_error_function = "FANN_ERRORFUNC_LINEAR"
204  mlpFANNCombiner.m_training_method = "FANN_TRAIN_RPROP"
205  mlpFANNCombiner.m_validation_fraction = 0.5
206  mlpFANNCombiner.m_random_seeds = 10
207  mlpFANNCombiner.m_test_rate = 500
208  mlpFANNCombiner.m_number_of_threads = 8
209  mlpFANNCombiner.m_scale_features = True
210  mlpFANNCombiner.m_scale_target = False
211  # mlpFANNCombiner.m_scale_target = True
212  return mlpFANNCombiner
213 
214 
215 # SignalFraction: FBDT feature
216 # For smooth output set to -1, this will break the calibration.
217 # For correct calibration set to -2, leads to peaky combiner output.
218 signalFraction = -2
219 
220 # Maximal number of events to train each method
221 maxEventsNumber = 0 # 0 takes all the sampled events. The number in the past was 500000
222 
223 
224 FTCategoryParameters = collections.namedtuple('FTCategoryParameters',
225  ['particleList', 'trackName', 'eventName', 'variableName', 'code'])
226 # Definition of all available categories, 'standard category name':
227 # ['ParticleList', 'trackLevel category name', 'eventLevel category name',
228 # 'combinerLevel variable name', 'category code']
229 AvailableCategories = {
230  'Electron':
231  FTCategoryParameters('e+:inRoe', 'Electron', 'Electron',
232  'QpOf(e+:inRoe, isRightCategory(Electron), isRightCategory(Electron))',
233  0),
234  'IntermediateElectron':
235  FTCategoryParameters('e+:inRoe', 'IntermediateElectron', 'IntermediateElectron',
236  'QpOf(e+:inRoe, isRightCategory(IntermediateElectron), isRightCategory(IntermediateElectron))',
237  1),
238  'Muon':
239  FTCategoryParameters('mu+:inRoe', 'Muon', 'Muon',
240  'QpOf(mu+:inRoe, isRightCategory(Muon), isRightCategory(Muon))',
241  2),
242  'IntermediateMuon':
243  FTCategoryParameters('mu+:inRoe', 'IntermediateMuon', 'IntermediateMuon',
244  'QpOf(mu+:inRoe, isRightCategory(IntermediateMuon), isRightCategory(IntermediateMuon))',
245  3),
246  'KinLepton':
247  FTCategoryParameters('mu+:inRoe', 'KinLepton', 'KinLepton',
248  'QpOf(mu+:inRoe, isRightCategory(KinLepton), isRightCategory(KinLepton))',
249  4),
250  'IntermediateKinLepton':
251  FTCategoryParameters('mu+:inRoe', 'IntermediateKinLepton', 'IntermediateKinLepton',
252  'QpOf(mu+:inRoe, isRightCategory(IntermediateKinLepton), isRightCategory(IntermediateKinLepton))',
253  5),
254  'Kaon':
255  FTCategoryParameters('K+:inRoe', 'Kaon', 'Kaon',
256  'weightedQpOf(K+:inRoe, isRightCategory(Kaon), isRightCategory(Kaon))',
257  6),
258  'SlowPion':
259  FTCategoryParameters('pi+:inRoe', 'SlowPion', 'SlowPion',
260  'QpOf(pi+:inRoe, isRightCategory(SlowPion), isRightCategory(SlowPion))',
261  7),
262  'FastHadron':
263  FTCategoryParameters('pi+:inRoe', 'FastHadron', 'FastHadron',
264  'QpOf(pi+:inRoe, isRightCategory(FastHadron), isRightCategory(FastHadron))',
265  8),
266  'Lambda':
267  FTCategoryParameters('Lambda0:inRoe', 'Lambda', 'Lambda',
268  'weightedQpOf(Lambda0:inRoe, isRightCategory(Lambda), isRightCategory(Lambda))',
269  9),
270  'FSC':
271  FTCategoryParameters('pi+:inRoe', 'SlowPion', 'FSC',
272  'QpOf(pi+:inRoe, isRightCategory(FSC), isRightCategory(SlowPion))',
273  10),
274  'MaximumPstar':
275  FTCategoryParameters('pi+:inRoe', 'MaximumPstar', 'MaximumPstar',
276  'QpOf(pi+:inRoe, isRightCategory(MaximumPstar), isRightCategory(MaximumPstar))',
277  11),
278  'KaonPion':
279  FTCategoryParameters('K+:inRoe', 'Kaon', 'KaonPion',
280  'QpOf(K+:inRoe, isRightCategory(KaonPion), isRightCategory(Kaon))',
281  12),
282 }
283 
284 
285 def getTrainingVariables(category=None):
286  """
287  Helper function to get training variables.
288 
289  NOTE: This function is not called the Expert mode. It is not necessary to be consistent with variables list of weight files.
290  """
291 
292  KId = {'Belle': 'ifNANgiveX(atcPIDBelle(3,2), 0.5)', 'Belle2': 'kaonID'}
293  muId = {'Belle': 'muIDBelle', 'Belle2': 'muonID'}
294  eId = {'Belle': 'eIDBelle', 'Belle2': 'electronID'}
295 
296  variables = []
297  if category == 'Electron' or category == 'IntermediateElectron':
298  variables = ['useCMSFrame(p)',
299  'useCMSFrame(pt)',
300  'p',
301  'pt',
302  'cosTheta',
303  eId[getBelleOrBelle2()],
304  'eid_TOP',
305  'eid_ARICH',
306  'eid_ECL',
307  'BtagToWBosonVariables(recoilMassSqrd)',
308  'BtagToWBosonVariables(pMissCMS)',
309  'BtagToWBosonVariables(cosThetaMissCMS)',
310  'BtagToWBosonVariables(EW90)',
311  'cosTPTO',
312  'chiProb',
313  ]
314  if getBelleOrBelle2() == "Belle":
315  variables.append('eid_dEdx')
316  variables.append('ImpactXY')
317  variables.append('distance')
318 
319  elif category == 'Muon' or category == 'IntermediateMuon':
320  variables = ['useCMSFrame(p)',
321  'useCMSFrame(pt)',
322  'p',
323  'pt',
324  'cosTheta',
325  muId[getBelleOrBelle2()],
326  'muid_TOP',
327  'muid_ARICH',
328  'muid_KLM',
329  'BtagToWBosonVariables(recoilMassSqrd)',
330  'BtagToWBosonVariables(pMissCMS)',
331  'BtagToWBosonVariables(cosThetaMissCMS)',
332  'BtagToWBosonVariables(EW90)',
333  'cosTPTO',
334  ]
335  if getBelleOrBelle2() == "Belle":
336  variables.append('muid_dEdx')
337  variables.append('ImpactXY')
338  variables.append('distance')
339  variables.append('chiProb')
340 
341  elif category == 'KinLepton' or category == 'IntermediateKinLepton':
342  variables = ['useCMSFrame(p)',
343  'useCMSFrame(pt)',
344  'p',
345  'pt',
346  'cosTheta',
347  muId[getBelleOrBelle2()],
348  'muid_TOP',
349  'muid_ARICH',
350  'muid_KLM',
351  eId[getBelleOrBelle2()],
352  'eid_TOP',
353  'eid_ARICH',
354  'eid_ECL',
355  'BtagToWBosonVariables(recoilMassSqrd)',
356  'BtagToWBosonVariables(pMissCMS)',
357  'BtagToWBosonVariables(cosThetaMissCMS)',
358  'BtagToWBosonVariables(EW90)',
359  'cosTPTO',
360  ]
361  if getBelleOrBelle2() == "Belle":
362  variables.append('eid_dEdx')
363  variables.append('muid_dEdx')
364  variables.append('ImpactXY')
365  variables.append('distance')
366  variables.append('chiProb')
367 
368  elif category == 'Kaon':
369  variables = ['useCMSFrame(p)',
370  'useCMSFrame(pt)',
371  'p',
372  'pt',
373  'cosTheta',
374  KId[getBelleOrBelle2()],
375  'Kid_dEdx',
376  'Kid_TOP',
377  'Kid_ARICH',
378  'NumberOfKShortsInRoe',
379  'ptTracksRoe',
380  'BtagToWBosonVariables(recoilMassSqrd)',
381  'BtagToWBosonVariables(pMissCMS)',
382  'BtagToWBosonVariables(cosThetaMissCMS)',
383  'BtagToWBosonVariables(EW90)',
384  'cosTPTO',
385  'chiProb',
386  ]
387  if getBelleOrBelle2() == "Belle":
388  variables.append('ImpactXY')
389  variables.append('distance')
390 
391  elif category == 'SlowPion':
392  variables = ['useCMSFrame(p)',
393  'useCMSFrame(pt)',
394  'cosTheta',
395  'p',
396  'pt',
397  'pionID',
398  'piid_TOP',
399  'piid_ARICH',
400  'pi_vs_edEdxid',
401  KId[getBelleOrBelle2()],
402  'Kid_dEdx',
403  'Kid_TOP',
404  'Kid_ARICH',
405  'NumberOfKShortsInRoe',
406  'ptTracksRoe',
407  eId[getBelleOrBelle2()],
408  'BtagToWBosonVariables(recoilMassSqrd)',
409  'BtagToWBosonVariables(EW90)',
410  'BtagToWBosonVariables(cosThetaMissCMS)',
411  'BtagToWBosonVariables(pMissCMS)',
412  'cosTPTO',
413  ]
414  if getBelleOrBelle2() == "Belle":
415  variables.append('piid_dEdx')
416  variables.append('ImpactXY')
417  variables.append('distance')
418  variables.append('chiProb')
419 
420  elif category == 'FastHadron':
421  variables = ['useCMSFrame(p)',
422  'useCMSFrame(pt)',
423  'cosTheta',
424  'p',
425  'pt',
426  'pionID',
427  'piid_dEdx',
428  'piid_TOP',
429  'piid_ARICH',
430  'pi_vs_edEdxid',
431  KId[getBelleOrBelle2()],
432  'Kid_dEdx',
433  'Kid_TOP',
434  'Kid_ARICH',
435  'NumberOfKShortsInRoe',
436  'ptTracksRoe',
437  eId[getBelleOrBelle2()],
438  'BtagToWBosonVariables(recoilMassSqrd)',
439  'BtagToWBosonVariables(EW90)',
440  'BtagToWBosonVariables(cosThetaMissCMS)',
441  'cosTPTO',
442  ]
443  if getBelleOrBelle2() == "Belle":
444  variables.append('BtagToWBosonVariables(pMissCMS)')
445  variables.append('ImpactXY')
446  variables.append('distance')
447  variables.append('chiProb')
448 
449  elif category == 'Lambda':
450  variables = ['lambdaFlavor',
451  'NumberOfKShortsInRoe',
452  'M',
453  'cosAngleBetweenMomentumAndVertexVector',
454  'lambdaZError',
455  'daughter(0,p)',
456  'daughter(0,useCMSFrame(p))',
457  'daughter(1,p)',
458  'daughter(1,useCMSFrame(p))',
459  'useCMSFrame(p)',
460  'p',
461  'chiProb',
462  ]
463  if getBelleOrBelle2() == "Belle2":
464  variables.append('daughter(1,protonID)') # protonID always 0 in B2BII check in future
465  variables.append('daughter(0,pionID)') # not very powerful in B2BII
466  else:
467  variables.append('distance')
468 
469  elif category == 'MaximumPstar':
470  variables = ['useCMSFrame(p)',
471  'useCMSFrame(pt)',
472  'p',
473  'pt',
474  'cosTPTO',
475  ]
476  if getBelleOrBelle2() == "Belle2":
477  variables.append('ImpactXY')
478  variables.append('distance')
479 
480  elif category == 'FSC':
481  variables = ['useCMSFrame(p)',
482  'cosTPTO',
483  KId[getBelleOrBelle2()],
484  'FSCVariables(pFastCMS)',
485  'FSCVariables(cosSlowFast)',
486  'FSCVariables(cosTPTOFast)',
487  'FSCVariables(SlowFastHaveOpositeCharges)',
488  ]
489  elif category == 'KaonPion':
490  variables = ['extraInfo(isRightCategory(Kaon))',
491  'HighestProbInCat(pi+:inRoe, isRightCategory(SlowPion))',
492  'KaonPionVariables(cosKaonPion)',
493  'KaonPionVariables(HaveOpositeCharges)',
494  KId[getBelleOrBelle2()]
495  ]
496 
497  return variables
498 
499 
500 def FillParticleLists(maskName='all', categories=None, path=None):
501  """
502  Fills the particle Lists for all categories.
503  """
504 
505  from vertex import kFit
506  readyParticleLists = []
507 
508  if categories is None:
509  categories = []
510 
511  trackCut = 'isInRestOfEvent > 0.5 and passesROEMask(' + maskName + ') > 0.5 and p >= 0'
512 
513  for category in categories:
514  particleList = AvailableCategories[category].particleList
515 
516  if particleList in readyParticleLists:
517  continue
518 
519  # Select particles in ROE for different categories according to mass hypothesis.
520  if particleList == 'Lambda0:inRoe':
521  if 'pi+:inRoe' not in readyParticleLists:
522  ma.fillParticleList('pi+:inRoe', trackCut, path=path)
523  readyParticleLists.append('pi+:inRoe')
524 
525  ma.fillParticleList('p+:inRoe', trackCut, path=path)
526  ma.reconstructDecay(particleList + ' -> pi-:inRoe p+:inRoe', '1.00<=M<=1.23', False, path=path)
527  kFit(particleList, 0.01, path=path)
528  ma.matchMCTruth(particleList, path=path)
529  readyParticleLists.append(particleList)
530 
531  else:
532  # Filling particle list for actual category
533  ma.fillParticleList(particleList, trackCut, path=path)
534  readyParticleLists.append(particleList)
535 
536  # Additional particleLists for K_S0
537  if getBelleOrBelle2() == 'Belle':
538  ma.cutAndCopyList('K_S0:inRoe', 'K_S0:mdst', 'extraInfo(ksnbStandard) == 1 and isInRestOfEvent == 1', path=path)
539  else:
540  if 'pi+:inRoe' not in readyParticleLists:
541  ma.fillParticleList('pi+:inRoe', trackCut, path=path)
542  ma.reconstructDecay('K_S0:inRoe -> pi+:inRoe pi-:inRoe', '0.40<=M<=0.60', False, path=path)
543  kFit('K_S0:inRoe', 0.01, path=path)
544 
545  # Apply BDT-based LID
546  if getBelleOrBelle2() == 'Belle2':
547  default_list_for_lid_BDT = ['e+:inRoe', 'mu+:inRoe']
548  list_for_lid_BDT = []
549 
550  for particleList in default_list_for_lid_BDT:
551  if particleList in readyParticleLists:
552  list_for_lid_BDT.append(particleList)
553 
554  if list_for_lid_BDT: # empty check
555  ma.applyChargedPidMVA(particleLists=list_for_lid_BDT, path=path,
556  trainingMode=0, # binary
557  binaryHypoPDGCodes=(11, 211)) # e vs pi
558  ma.applyChargedPidMVA(particleLists=list_for_lid_BDT, path=path,
559  trainingMode=0, # binary
560  binaryHypoPDGCodes=(13, 211)) # mu vs pi
561  ma.applyChargedPidMVA(particleLists=list_for_lid_BDT, path=path,
562  trainingMode=1) # Multiclass
563 
564 
565 def eventLevel(mode='Expert', weightFiles='B2JpsiKs_mu', categories=None, path=None):
566  """
567  Samples data for training or tests all categories all categories at event level.
568  """
569 
570  from basf2 import create_path
571  from basf2 import register_module
572 
573  B2INFO('EVENT LEVEL')
574 
575  ReadyMethods = 0
576 
577  # Each category has its own Path in order to be skipped if the corresponding particle list is empty
578  identifiersExtraInfosDict = dict()
579  identifiersExtraInfosKaonPion = []
580 
581  if categories is None:
582  categories = []
583 
584  for category in categories:
585  particleList = AvailableCategories[category].particleList
586 
587  methodPrefixEventLevel = "FlavorTagger_" + getBelleOrBelle2() + "_" + weightFiles + 'EventLevel' + category + 'FBDT'
588  identifierEventLevel = methodPrefixEventLevel
589  targetVariable = 'isRightCategory(' + category + ')'
590  extraInfoName = targetVariable
591 
592  if mode == 'Expert':
593 
594  if downloadFlag or useOnlyLocalFlag:
595  identifierEventLevel = filesDirectory + '/' + methodPrefixEventLevel + '_1.root'
596 
597  if downloadFlag:
598  if not os.path.isfile(identifierEventLevel):
599  basf2_mva.download(methodPrefixEventLevel, identifierEventLevel)
600  if not os.path.isfile(identifierEventLevel):
601  B2FATAL('Flavor Tagger: Weight file ' + identifierEventLevel +
602  ' was not downloaded from Database. Please check the buildOrRevision name. Stopped')
603 
604  if useOnlyLocalFlag:
605  if not os.path.isfile(identifierEventLevel):
606  B2FATAL('Flavor Tagger: ' + particleList + ' Eventlevel was not trained. Weight file ' +
607  identifierEventLevel + ' was not found. Stopped')
608 
609  B2INFO('flavorTagger: MVAExpert ' + methodPrefixEventLevel + ' ready.')
610 
611  elif mode == 'Sampler':
612 
613  identifierEventLevel = filesDirectory + '/' + methodPrefixEventLevel + '_1.root'
614  if os.path.isfile(identifierEventLevel):
615  B2INFO('flavorTagger: MVAExpert ' + methodPrefixEventLevel + ' ready.')
616 
617  if 'KaonPion' in categories:
618  methodPrefixEventLevelKaonPion = "FlavorTagger_" + getBelleOrBelle2() + \
619  "_" + weightFiles + 'EventLevelKaonPionFBDT'
620  identifierEventLevelKaonPion = filesDirectory + '/' + methodPrefixEventLevelKaonPion + '_1.root'
621  if not os.path.isfile(identifierEventLevelKaonPion):
622  # Slow Pion and Kaon categories are used if Kaon-Pion is lacking for
623  # sampling. The others are not needed and skipped
624  if category != "SlowPion" and category != "Kaon":
625  continue
626 
627  if mode == 'Expert' or (mode == 'Sampler' and os.path.isfile(identifierEventLevel)):
628 
629  B2INFO('flavorTagger: Applying MVAExpert ' + methodPrefixEventLevel + '.')
630 
631  if category == 'KaonPion':
632  identifiersExtraInfosKaonPion.append((extraInfoName, identifierEventLevel))
633  elif particleList not in identifiersExtraInfosDict:
634  identifiersExtraInfosDict[particleList] = [(extraInfoName, identifierEventLevel)]
635  else:
636  identifiersExtraInfosDict[particleList].append((extraInfoName, identifierEventLevel))
637 
638  ReadyMethods += 1
639 
640  # Each category has its own Path in order to be skipped if the corresponding particle list is empty
641  for particleList in identifiersExtraInfosDict:
642  eventLevelPath = create_path()
643  SkipEmptyParticleList = register_module("SkimFilter")
644  SkipEmptyParticleList.set_name('SkimFilter_EventLevel_' + particleList)
645  SkipEmptyParticleList.param('particleLists', particleList)
646  SkipEmptyParticleList.if_true(eventLevelPath, basf2.AfterConditionPath.CONTINUE)
647  path.add_module(SkipEmptyParticleList)
648 
649  mvaMultipleExperts = register_module('MVAMultipleExperts')
650  mvaMultipleExperts.set_name('MVAMultipleExperts_EventLevel_' + particleList)
651  mvaMultipleExperts.param('listNames', [particleList])
652  mvaMultipleExperts.param('extraInfoNames', [row[0] for row in identifiersExtraInfosDict[particleList]])
653  mvaMultipleExperts.param('signalFraction', signalFraction)
654  mvaMultipleExperts.param('identifiers', [row[1] for row in identifiersExtraInfosDict[particleList]])
655  eventLevelPath.add_module(mvaMultipleExperts)
656 
657  if 'KaonPion' in categories and len(identifiersExtraInfosKaonPion) != 0:
658  eventLevelKaonPionPath = create_path()
659  SkipEmptyParticleList = register_module("SkimFilter")
660  SkipEmptyParticleList.set_name('SkimFilter_' + 'K+:inRoe')
661  SkipEmptyParticleList.param('particleLists', 'K+:inRoe')
662  SkipEmptyParticleList.if_true(eventLevelKaonPionPath, basf2.AfterConditionPath.CONTINUE)
663  path.add_module(SkipEmptyParticleList)
664 
665  mvaExpertKaonPion = register_module("MVAExpert")
666  mvaExpertKaonPion.set_name('MVAExpert_KaonPion_' + 'K+:inRoe')
667  mvaExpertKaonPion.param('listNames', ['K+:inRoe'])
668  mvaExpertKaonPion.param('extraInfoName', identifiersExtraInfosKaonPion[0][0])
669  mvaExpertKaonPion.param('signalFraction', signalFraction)
670  mvaExpertKaonPion.param('identifier', identifiersExtraInfosKaonPion[0][1])
671 
672  eventLevelKaonPionPath.add_module(mvaExpertKaonPion)
673 
674  if mode == 'Sampler':
675 
676  for category in categories:
677  particleList = AvailableCategories[category].particleList
678 
679  methodPrefixEventLevel = "FlavorTagger_" + getBelleOrBelle2() + "_" + weightFiles + 'EventLevel' + category + 'FBDT'
680  identifierEventLevel = filesDirectory + '/' + methodPrefixEventLevel + '_1.root'
681  targetVariable = 'isRightCategory(' + category + ')'
682 
683  if not os.path.isfile(identifierEventLevel):
684 
685  if category == 'KaonPion':
686  methodPrefixEventLevelSlowPion = "FlavorTagger_" + getBelleOrBelle2() + \
687  "_" + weightFiles + 'EventLevelSlowPionFBDT'
688  identifierEventLevelSlowPion = filesDirectory + '/' + methodPrefixEventLevelSlowPion + '_1.root'
689  if not os.path.isfile(identifierEventLevelSlowPion):
690  B2INFO("Flavor Tagger: event level weight file for the Slow Pion category is absent." +
691  "It is required to sample the training information for the KaonPion category." +
692  "An additional sampling step will be needed after the following training step.")
693  continue
694 
695  B2INFO('flavorTagger: file ' + filesDirectory + '/' +
696  methodPrefixEventLevel + "sampled" + fileId + '.root will be saved.')
697 
698  ma.applyCuts(particleList, 'isRightCategory(mcAssociated) > 0', path)
699  eventLevelpath = create_path()
700  SkipEmptyParticleList = register_module("SkimFilter")
701  SkipEmptyParticleList.set_name('SkimFilter_EventLevel' + category)
702  SkipEmptyParticleList.param('particleLists', particleList)
703  SkipEmptyParticleList.if_true(eventLevelpath, basf2.AfterConditionPath.CONTINUE)
704  path.add_module(SkipEmptyParticleList)
705 
706  ntuple = register_module('VariablesToNtuple')
707  ntuple.param('fileName', filesDirectory + '/' + methodPrefixEventLevel + "sampled" + fileId + ".root")
708  ntuple.param('treeName', methodPrefixEventLevel + "_tree")
709  variablesToBeSaved = getTrainingVariables(category) + [targetVariable, 'ancestorHasWhichFlavor',
710  'isSignal', 'mcPDG', 'mcErrors', 'genMotherPDG',
711  'nMCMatches', 'B0mcErrors']
712  if category != 'KaonPion' and category != 'FSC':
713  variablesToBeSaved = variablesToBeSaved + \
714  ['extraInfo(isRightTrack(' + category + '))',
715  'hasHighestProbInCat(' + particleList + ', isRightTrack(' + category + '))']
716  ntuple.param('variables', variablesToBeSaved)
717  ntuple.param('particleList', particleList)
718  eventLevelpath.add_module(ntuple)
719 
720  if ReadyMethods != len(categories):
721  return False
722  else:
723  return True
724 
725 
726 def eventLevelTeacher(weightFiles='B2JpsiKs_mu', categories=None):
727  """
728  Trains all categories at event level.
729  """
730 
731  B2INFO('EVENT LEVEL TEACHER')
732 
733  ReadyMethods = 0
734 
735  if categories is None:
736  categories = []
737 
738  for category in categories:
739  methodPrefixEventLevel = "FlavorTagger_" + getBelleOrBelle2() + "_" + weightFiles + 'EventLevel' + category + 'FBDT'
740  targetVariable = 'isRightCategory(' + category + ')'
741  weightFile = filesDirectory + '/' + methodPrefixEventLevel + "_1.root"
742 
743  if os.path.isfile(weightFile):
744  ReadyMethods += 1
745  continue
746 
747  sampledFilesList = glob.glob(filesDirectory + '/' + methodPrefixEventLevel + 'sampled*.root')
748  if len(sampledFilesList) == 0:
749  B2INFO('flavorTagger: eventLevelTeacher did not find any ' + methodPrefixEventLevel +
750  ".root" + ' file. Please run the flavorTagger in "Sampler" mode afterwards.')
751 
752  else:
753  B2INFO('flavorTagger: MVA Teacher training' + methodPrefixEventLevel + ' .')
754  trainingOptionsEventLevel = basf2_mva.GeneralOptions()
755  trainingOptionsEventLevel.m_datafiles = basf2_mva.vector(*sampledFilesList)
756  trainingOptionsEventLevel.m_treename = methodPrefixEventLevel + "_tree"
757  trainingOptionsEventLevel.m_identifier = weightFile
758  trainingOptionsEventLevel.m_variables = basf2_mva.vector(*getTrainingVariables(category))
759  trainingOptionsEventLevel.m_target_variable = targetVariable
760  trainingOptionsEventLevel.m_max_events = maxEventsNumber
761 
762  basf2_mva.teacher(trainingOptionsEventLevel, getFastBDTCategories())
763 
764  if uploadFlag:
765  basf2_mva.upload(weightFile, methodPrefixEventLevel)
766 
767  if ReadyMethods != len(categories):
768  return False
769  else:
770  return True
771 
772 
773 def combinerLevel(mode='Expert', weightFiles='B2JpsiKs_mu', categories=None,
774  variablesCombinerLevel=None, categoriesCombinationCode=None, path=None):
775  """
776  Samples the input data or tests the combiner according to the selected categories.
777  """
778 
779  B2INFO('COMBINER LEVEL')
780 
781  if categories is None:
782  categories = []
783  if variablesCombinerLevel is None:
784  variablesCombinerLevel = []
785 
786  B2INFO("Flavor Tagger: Required Combiner for Categories:")
787  for category in categories:
788  B2INFO(category)
789 
790  B2INFO("Flavor Tagger: which corresponds to a weight file with categories combination code " + categoriesCombinationCode)
791 
792  methodPrefixCombinerLevel = "FlavorTagger_" + getBelleOrBelle2() + "_" + weightFiles + 'Combiner' \
793  + categoriesCombinationCode
794 
795  if mode == 'Sampler':
796 
797  if os.path.isfile(filesDirectory + '/' + methodPrefixCombinerLevel + 'FBDT' + '_1.root') or \
798  os.path.isfile(filesDirectory + '/' + methodPrefixCombinerLevel + 'FANN' + '_1.root'):
799  B2FATAL('flavorTagger: File' + methodPrefixCombinerLevel + 'FBDT' + "_1.root" + ' or ' + methodPrefixCombinerLevel +
800  'FANN' + '_1.root found. Please run the "Expert" mode or delete the file if a new sampling is desired.')
801 
802  B2INFO('flavorTagger: Sampling Data on Combiner Level. File' +
803  methodPrefixCombinerLevel + ".root" + ' will be saved')
804 
805  ntuple = basf2.register_module('VariablesToNtuple')
806  ntuple.param('fileName', filesDirectory + '/' + methodPrefixCombinerLevel + "sampled" + fileId + ".root")
807  ntuple.param('treeName', methodPrefixCombinerLevel + 'FBDT' + "_tree")
808  ntuple.param('variables', variablesCombinerLevel + ['qrCombined'])
809  ntuple.param('particleList', "")
810  path.add_module(ntuple)
811 
812  if mode == 'Expert':
813 
814  # Check if weight files are ready
815  if TMVAfbdt:
816  identifierFBDT = methodPrefixCombinerLevel + 'FBDT'
817  if downloadFlag or useOnlyLocalFlag:
818  identifierFBDT = filesDirectory + '/' + methodPrefixCombinerLevel + 'FBDT' + '_1.root'
819 
820  if downloadFlag:
821  if not os.path.isfile(identifierFBDT):
822  basf2_mva.download(methodPrefixCombinerLevel + 'FBDT', identifierFBDT)
823  if not os.path.isfile(identifierFBDT):
824  B2FATAL('Flavor Tagger: Weight file ' + identifierFBDT +
825  ' was not downloaded from Database. Please check the buildOrRevision name. Stopped')
826 
827  if useOnlyLocalFlag:
828  if not os.path.isfile(identifierFBDT):
829  B2FATAL('flavorTagger: Combinerlevel FastBDT was not trained with this combination of categories.' +
830  ' Weight file ' + identifierFBDT + ' not found. Stopped')
831 
832  B2INFO('flavorTagger: Ready to be used with weightFile ' + methodPrefixCombinerLevel + 'FBDT' + '_1.root')
833 
834  if FANNmlp:
835  identifierFANN = methodPrefixCombinerLevel + 'FANN'
836  if downloadFlag or useOnlyLocalFlag:
837  identifierFANN = filesDirectory + '/' + methodPrefixCombinerLevel + 'FANN' + '_1.root'
838 
839  if downloadFlag:
840  if not os.path.isfile(identifierFANN):
841  basf2_mva.download(methodPrefixCombinerLevel + 'FANN', identifierFANN)
842  if not os.path.isfile(identifierFANN):
843  B2FATAL('Flavor Tagger: Weight file ' + identifierFANN +
844  ' was not downloaded from Database. Please check the buildOrRevision name. Stopped')
845  if useOnlyLocalFlag:
846  if not os.path.isfile(identifierFANN):
847  B2FATAL('flavorTagger: Combinerlevel FANNMLP was not trained with this combination of categories. ' +
848  ' Weight file ' + identifierFANN + ' not found. Stopped')
849 
850  B2INFO('flavorTagger: Ready to be used with weightFile ' + methodPrefixCombinerLevel + 'FANN' + '_1.root')
851 
852  # At this stage, all necessary weight files should be ready.
853  # Call MVAExpert or MVAMultipleExperts module.
854  if TMVAfbdt and not FANNmlp:
855  B2INFO('flavorTagger: Apply FBDTMethod ' + methodPrefixCombinerLevel + 'FBDT')
856  path.add_module('MVAExpert', listNames=[], extraInfoName='qrCombined' + 'FBDT', signalFraction=signalFraction,
857  identifier=identifierFBDT)
858 
859  if FANNmlp and not TMVAfbdt:
860  B2INFO('flavorTagger: Apply FANNMethod on combiner level')
861  path.add_module('MVAExpert', listNames=[], extraInfoName='qrCombined' + 'FANN', signalFraction=signalFraction,
862  identifier=identifierFANN)
863 
864  if FANNmlp and TMVAfbdt:
865  B2INFO('flavorTagger: Apply FANNMethod and FBDTMethod on combiner level')
866  mvaMultipleExperts = basf2.register_module('MVAMultipleExperts')
867  mvaMultipleExperts.set_name('MVAMultipleExperts_Combiners')
868  mvaMultipleExperts.param('listNames', [])
869  mvaMultipleExperts.param('extraInfoNames', ['qrCombined' + 'FBDT', 'qrCombined' + 'FANN'])
870  mvaMultipleExperts.param('signalFraction', signalFraction)
871  mvaMultipleExperts.param('identifiers', [identifierFBDT, identifierFANN])
872  path.add_module(mvaMultipleExperts)
873 
874 
875 def combinerLevelTeacher(weightFiles='B2JpsiKs_mu', variablesCombinerLevel=None,
876  categoriesCombinationCode=None):
877  """
878  Trains the combiner according to the selected categories.
879  """
880 
881  B2INFO('COMBINER LEVEL TEACHER')
882 
883  if variablesCombinerLevel is None:
884  variablesCombinerLevel = []
885 
886  methodPrefixCombinerLevel = "FlavorTagger_" + getBelleOrBelle2() + "_" + weightFiles + 'Combiner' \
887  + categoriesCombinationCode
888 
889  sampledFilesList = glob.glob(filesDirectory + '/' + methodPrefixCombinerLevel + 'sampled*.root')
890  if len(sampledFilesList) == 0:
891  B2FATAL('FlavorTagger: combinerLevelTeacher did not find any ' +
892  methodPrefixCombinerLevel + 'sampled*.root file. Please run the flavorTagger in "Sampler" mode.')
893 
894  if TMVAfbdt:
895 
896  if not os.path.isfile(filesDirectory + '/' + methodPrefixCombinerLevel + 'FBDT' + '_1.root'):
897 
898  B2INFO('flavorTagger: MVA Teacher training a FastBDT on Combiner Level')
899 
900  trainingOptionsCombinerLevel = basf2_mva.GeneralOptions()
901  trainingOptionsCombinerLevel.m_datafiles = basf2_mva.vector(*sampledFilesList)
902  trainingOptionsCombinerLevel.m_treename = methodPrefixCombinerLevel + 'FBDT' + "_tree"
903  trainingOptionsCombinerLevel.m_identifier = filesDirectory + '/' + methodPrefixCombinerLevel + 'FBDT' + "_1.root"
904  trainingOptionsCombinerLevel.m_variables = basf2_mva.vector(*variablesCombinerLevel)
905  trainingOptionsCombinerLevel.m_target_variable = 'qrCombined'
906  trainingOptionsCombinerLevel.m_max_events = maxEventsNumber
907 
908  basf2_mva.teacher(trainingOptionsCombinerLevel, getFastBDTCombiner())
909 
910  if uploadFlag:
911  basf2_mva.upload(filesDirectory + '/' + methodPrefixCombinerLevel +
912  'FBDT' + "_1.root", methodPrefixCombinerLevel + 'FBDT')
913 
914  elif FANNmlp and not os.path.isfile(filesDirectory + '/' + methodPrefixCombinerLevel + 'FANN' + '_1.root'):
915 
916  B2INFO('flavorTagger: Combinerlevel FBDT was already trained with this combination of categories. Weight file ' +
917  methodPrefixCombinerLevel + 'FBDT' + '_1.root has been found.')
918 
919  else:
920  B2FATAL('flavorTagger: Combinerlevel was already trained with this combination of categories. Weight files ' +
921  methodPrefixCombinerLevel + 'FBDT' + '_1.root and ' +
922  methodPrefixCombinerLevel + 'FANN' + '_1.root has been found. Please use the "Expert" mode')
923 
924  if FANNmlp:
925 
926  if not os.path.isfile(filesDirectory + '/' + methodPrefixCombinerLevel + 'FANN' + '_1.root'):
927 
928  B2INFO('flavorTagger: MVA Teacher training a FANN MLP on Combiner Level')
929 
930  trainingOptionsCombinerLevel = basf2_mva.GeneralOptions()
931  trainingOptionsCombinerLevel.m_datafiles = basf2_mva.vector(*sampledFilesList)
932  trainingOptionsCombinerLevel.m_treename = methodPrefixCombinerLevel + 'FBDT' + "_tree"
933  trainingOptionsCombinerLevel.m_identifier = filesDirectory + '/' + methodPrefixCombinerLevel + 'FANN' + "_1.root"
934  trainingOptionsCombinerLevel.m_variables = basf2_mva.vector(*variablesCombinerLevel)
935  trainingOptionsCombinerLevel.m_target_variable = 'qrCombined'
936  trainingOptionsCombinerLevel.m_max_events = maxEventsNumber
937 
938  basf2_mva.teacher(trainingOptionsCombinerLevel, getMlpFANNCombiner())
939 
940  if uploadFlag:
941  basf2_mva.upload(filesDirectory + '/' + methodPrefixCombinerLevel +
942  'FANN' + "_1.root", methodPrefixCombinerLevel + 'FANN')
943 
944  elif TMVAfbdt and not os.path.isfile(filesDirectory + '/' + methodPrefixCombinerLevel + 'FBDT' + '_1.root'):
945 
946  B2INFO('flavorTagger: Combinerlevel FBDT was already trained with this combination of categories. Weight file ' +
947  methodPrefixCombinerLevel + 'FANN' + '_1.config has been found.')
948 
949  else:
950  B2FATAL('flavorTagger: Combinerlevel was already trained with this combination of categories. Weight files ' +
951  methodPrefixCombinerLevel + 'FBDT' + '_1.root and ' +
952  methodPrefixCombinerLevel + 'FANN' + '_1.root has been found. Please use the "Expert" mode')
953 
954 
955 def getEventLevelParticleLists(categories=None):
956 
957  if categories is None:
958  categories = []
959 
960  eventLevelParticleLists = []
961 
962  for category in categories:
963  ftCategory = AvailableCategories[category]
964  event_tuple = (ftCategory.particleList, ftCategory.eventName, ftCategory.variableName)
965 
966  if event_tuple not in eventLevelParticleLists:
967  eventLevelParticleLists.append(event_tuple)
968  else:
969  B2FATAL('Flavor Tagger: ' + category + ' has been already given')
970 
971  return eventLevelParticleLists
972 
973 
974 def flavorTagger(
975  particleLists=None,
976  mode='Expert',
977  weightFiles='B2nunubarBGx1',
978  workingDirectory='.',
979  combinerMethods=['TMVA-FBDT'],
980  categories=[
981  'Electron',
982  'IntermediateElectron',
983  'Muon',
984  'IntermediateMuon',
985  'KinLepton',
986  'IntermediateKinLepton',
987  'Kaon',
988  'SlowPion',
989  'FastHadron',
990  'Lambda',
991  'FSC',
992  'MaximumPstar',
993  'KaonPion'],
994  maskName='FTDefaultMask',
995  saveCategoriesInfo=True,
996  useOnlyLocalWeightFiles=False,
997  downloadFromDatabaseIfNotFound=False,
998  uploadToDatabaseAfterTraining=False,
999  samplerFileId='',
1000  prefix='MC15ri_light-2207-bengal_0',
1001  path=None,
1002 ):
1003  """
1004  Defines the whole flavor tagging process for each selected Rest of Event (ROE) built in the steering file.
1005  The flavor is predicted by Multivariate Methods trained with Variables and MetaVariables which use
1006  Tracks, ECL- and KLMClusters from the corresponding RestOfEvent dataobject.
1007  This module can be used to sample the training information, to train and/or to test the flavorTagger.
1008 
1009  @param particleLists The ROEs for flavor tagging are selected from the given particle lists.
1010  @param mode The available modes are
1011  ``Expert`` (default), ``Sampler``, and ``Teacher``. In the ``Expert`` mode
1012  Flavor Tagging is applied to the analysis,. In the ``Sampler`` mode you save
1013  save the variables for training. In the ``Teacher`` mode the FlavorTagger is
1014  trained, for this step you do not reconstruct any particle or do any analysis,
1015  you just run the flavorTagger alone.
1016  @param weightFiles Weight files name. Default=
1017  ``B2nunubarBGx1`` (official weight files). If the user self
1018  wants to train the FlavorTagger, the weightfiles name should correspond to the
1019  analysed CP channel in order to avoid confusions. The default name
1020  ``B2nunubarBGx1`` corresponds to
1021  :math:`B^0_{\\rm sig}\\to \\nu \\overline{\\nu}`.
1022  and ``B2JpsiKs_muBGx1`` to
1023  :math:`B^0_{\\rm sig}\\to J/\\psi (\\to \\mu^+ \\mu^-) K_s (\\to \\pi^+ \\pi^-)`.
1024  BGx1 stays for events simulated with background.
1025  @param workingDirectory Path to the directory containing the FlavorTagging/ folder.
1026  @param combinerMethods MVAs for the combiner: ``TMVA-FBDT` (default).
1027  ``FANN-MLP`` is available only with ``prefix=''`` (MC13 weight files).
1028  @param categories Categories used for flavor tagging. By default all are used.
1029  @param maskName Gets ROE particles from a specified ROE mask.
1030  ``FTDefaultMask`` (default): tentative mask definition that will be created
1031  automatically. The definition is as follows:
1032 
1033  - Track (pion): thetaInCDCAcceptance and dr<1 and abs(dz)<3
1034  - ECL-cluster (gamma): thetaInCDCAcceptance and clusterNHits>1.5 and \
1035  [[clusterReg==1 and E>0.08] or [clusterReg==2 and E>0.03] or \
1036  [clusterReg==3 and E>0.06]] \
1037  (Same as gamma:pi0eff30_May2020 and gamma:pi0eff40_May2020)
1038 
1039  ``all``: all ROE particles are used.
1040  Or one can give any mask name defined before calling this function.
1041  @param saveCategoriesInfo Sets to save information of individual categories.
1042  @param useOnlyLocalWeightFiles [Expert] Uses only locally saved weight files.
1043  @param downloadFromDatabaseIfNotFound [Expert] Weight files are downloaded from
1044  the conditions database if not available in workingDirectory.
1045  @param uploadToDatabaseAfterTraining [Expert] For librarians only: uploads weight files to localdb after training.
1046  @param samplerFileId Identifier to paralellize
1047  sampling. Only used in ``Sampler`` mode. If you are training by yourself and
1048  want to parallelize the sampling, you can run several sampling scripts in
1049  parallel. By changing this parameter you will not overwrite an older sample.
1050  @param prefix Prefix of weight files.
1051  ``MC15ri_light-2207-bengal_0`` (default): Weight files trained for MC15ri samples.
1052  ``''``: Weight files trained for MC13 samples.
1053  @param path Modules are added to this path
1054 
1055  """
1056 
1057  if (not isinstance(particleLists, list)):
1058  particleLists = [particleLists] # in case user inputs a particle list as string
1059 
1060  if mode != 'Sampler' and mode != 'Teacher' and mode != 'Expert':
1061  B2FATAL('flavorTagger: Wrong mode given: The available modes are "Sampler", "Teacher" or "Expert"')
1062 
1063  if len(categories) != len(set(categories)):
1064  dup = [cat for cat in set(categories) if categories.count(cat) > 1]
1065  B2WARNING('Flavor Tagger: There are duplicate elements in the given categories list. '
1066  << 'The following duplicate elements are removed; ' << ', '.join(dup))
1067  categories = list(set(categories))
1068 
1069  if len(categories) < 2:
1070  B2FATAL('Flavor Tagger: Invalid amount of categories. At least two are needed.')
1071  B2FATAL(
1072  'Flavor Tagger: Possible categories are "Electron", "IntermediateElectron", "Muon", "IntermediateMuon", '
1073  '"KinLepton", "IntermediateKinLepton", "Kaon", "SlowPion", "FastHadron",'
1074  '"Lambda", "FSC", "MaximumPstar" or "KaonPion" ')
1075 
1076  for category in categories:
1077  if category not in AvailableCategories:
1078  B2FATAL('Flavor Tagger: ' + category + ' is not a valid category name given')
1079  B2FATAL('Flavor Tagger: Available categories are "Electron", "IntermediateElectron", '
1080  '"Muon", "IntermediateMuon", "KinLepton", "IntermediateKinLepton", "Kaon", "SlowPion", "FastHadron", '
1081  '"Lambda", "FSC", "MaximumPstar" or "KaonPion" ')
1082 
1083  # Directory where the weights of the trained Methods are saved
1084  # workingDirectory = os.environ['BELLE2_LOCAL_DIR'] + '/analysis/data'
1085 
1086  basf2.find_file(workingDirectory)
1087 
1088  global filesDirectory
1089  filesDirectory = workingDirectory + '/FlavorTagging/TrainedMethods'
1090 
1091  if mode == 'Sampler' or (mode == 'Expert' and downloadFromDatabaseIfNotFound):
1092  if not basf2.find_file(workingDirectory + '/FlavorTagging', silent=True):
1093  os.mkdir(workingDirectory + '/FlavorTagging')
1094  os.mkdir(workingDirectory + '/FlavorTagging/TrainedMethods')
1095  elif not basf2.find_file(workingDirectory + '/FlavorTagging/TrainedMethods', silent=True):
1096  os.mkdir(workingDirectory + '/FlavorTagging/TrainedMethods')
1097  filesDirectory = workingDirectory + '/FlavorTagging/TrainedMethods'
1098 
1099  if len(combinerMethods) < 1 or len(combinerMethods) > 2:
1100  B2FATAL('flavorTagger: Invalid list of combinerMethods. The available methods are "TMVA-FBDT" and "FANN-MLP"')
1101 
1102  global FANNmlp
1103  global TMVAfbdt
1104 
1105  FANNmlp = False
1106  TMVAfbdt = False
1107 
1108  for method in combinerMethods:
1109  if method == 'TMVA-FBDT':
1110  TMVAfbdt = True
1111  elif method == 'FANN-MLP':
1112  FANNmlp = True
1113  else:
1114  B2FATAL('flavorTagger: Invalid list of combinerMethods. The available methods are "TMVA-FBDT" and "FANN-MLP"')
1115 
1116  global fileId
1117  fileId = samplerFileId
1118 
1119  global useOnlyLocalFlag
1120  useOnlyLocalFlag = useOnlyLocalWeightFiles
1121 
1122  B2INFO('*** FLAVOR TAGGING ***')
1123  B2INFO(' ')
1124  B2INFO(' Working directory is: ' + filesDirectory)
1125  B2INFO(' ')
1126 
1127  setInteractionWithDatabase(downloadFromDatabaseIfNotFound, uploadToDatabaseAfterTraining)
1128 
1129  if prefix == '':
1130  set_FlavorTagger_pid_aliases_legacy()
1131  else:
1132  set_FlavorTagger_pid_aliases()
1133 
1134  setInputVariablesWithMask()
1135  if prefix != '':
1136  weightFiles = prefix + '_' + weightFiles
1137 
1138  # Create configuration lists and code-name for given category's list
1139  trackLevelParticleLists = []
1140  eventLevelParticleLists = []
1141  variablesCombinerLevel = []
1142  categoriesCombination = []
1143  categoriesCombinationCode = 'CatCode'
1144  for category in categories:
1145  ftCategory = AvailableCategories[category]
1146 
1147  track_tuple = (ftCategory.particleList, ftCategory.trackName)
1148  event_tuple = (ftCategory.particleList, ftCategory.eventName, ftCategory.variableName)
1149 
1150  if track_tuple not in trackLevelParticleLists and category != 'MaximumPstar':
1151  trackLevelParticleLists.append(track_tuple)
1152 
1153  if event_tuple not in eventLevelParticleLists:
1154  eventLevelParticleLists.append(event_tuple)
1155  variablesCombinerLevel.append(ftCategory.variableName)
1156  categoriesCombination.append(ftCategory.code)
1157  else:
1158  B2FATAL('Flavor Tagger: ' + category + ' has been already given')
1159 
1160  for code in sorted(categoriesCombination):
1161  categoriesCombinationCode = categoriesCombinationCode + '%02d' % code
1162 
1163  # Create default ROE-mask
1164  if maskName == 'FTDefaultMask':
1165  FTDefaultMask = (
1166  'FTDefaultMask',
1167  'thetaInCDCAcceptance and dr<1 and abs(dz)<3',
1168  'thetaInCDCAcceptance and clusterNHits>1.5 and [[E>0.08 and clusterReg==1] or [E>0.03 and clusterReg==2] or \
1169  [E>0.06 and clusterReg==3]]')
1170  for name in particleLists:
1171  ma.appendROEMasks(list_name=name, mask_tuples=[FTDefaultMask], path=path)
1172 
1173  # Start ROE-routine
1174  roe_path = basf2.create_path()
1175  deadEndPath = basf2.create_path()
1176 
1177  if mode == 'Sampler':
1178  # Events containing ROE without B-Meson (but not empty) are discarded for training
1179  ma.signalSideParticleListsFilter(
1180  particleLists,
1181  'nROE_Charged(' + maskName + ', 0) > 0 and abs(qrCombined) == 1',
1182  roe_path,
1183  deadEndPath)
1184 
1185  FillParticleLists(maskName, categories, roe_path)
1186 
1187  if eventLevel(mode, weightFiles, categories, roe_path):
1188  combinerLevel(mode, weightFiles, categories, variablesCombinerLevel, categoriesCombinationCode, roe_path)
1189 
1190  path.for_each('RestOfEvent', 'RestOfEvents', roe_path)
1191 
1192  elif mode == 'Expert':
1193  # If trigger returns 1 jump into empty path skipping further modules in roe_path
1194  # run filter with no cut first to get rid of ROEs that are missing the mask of the signal particle
1195  ma.signalSideParticleListsFilter(particleLists, '', roe_path, deadEndPath)
1196  ma.signalSideParticleListsFilter(particleLists, 'nROE_Charged(' + maskName + ', 0) > 0', roe_path, deadEndPath)
1197 
1198  # Initialization of flavorTaggerInfo dataObject needs to be done in the main path
1199  flavorTaggerInfoBuilder = basf2.register_module('FlavorTaggerInfoBuilder')
1200  path.add_module(flavorTaggerInfoBuilder)
1201 
1202  FillParticleLists(maskName, categories, roe_path)
1203 
1204  if eventLevel(mode, weightFiles, categories, roe_path):
1205  combinerLevel(mode, weightFiles, categories, variablesCombinerLevel, categoriesCombinationCode, roe_path)
1206 
1207  flavorTaggerInfoFiller = basf2.register_module('FlavorTaggerInfoFiller')
1208  flavorTaggerInfoFiller.param('trackLevelParticleLists', trackLevelParticleLists)
1209  flavorTaggerInfoFiller.param('eventLevelParticleLists', eventLevelParticleLists)
1210  flavorTaggerInfoFiller.param('TMVAfbdt', TMVAfbdt)
1211  flavorTaggerInfoFiller.param('FANNmlp', FANNmlp)
1212  flavorTaggerInfoFiller.param('qpCategories', saveCategoriesInfo)
1213  flavorTaggerInfoFiller.param('istrueCategories', saveCategoriesInfo)
1214  flavorTaggerInfoFiller.param('targetProb', False)
1215  flavorTaggerInfoFiller.param('trackPointers', False)
1216  roe_path.add_module(flavorTaggerInfoFiller) # Add FlavorTag Info filler to roe_path
1217  add_default_FlavorTagger_aliases()
1218 
1219  path.for_each('RestOfEvent', 'RestOfEvents', roe_path)
1220 
1221  elif mode == 'Teacher':
1222  if eventLevelTeacher(weightFiles, categories):
1223  combinerLevelTeacher(weightFiles, variablesCombinerLevel, categoriesCombinationCode)
1224 
1225 
1226 if __name__ == '__main__':
1227 
1228  desc_list = []
1229 
1230  function = globals()["flavorTagger"]
1231  signature = inspect.formatargspec(*inspect.getfullargspec(function))
1232  desc_list.append((function.__name__, signature + '\n' + function.__doc__))
1233 
1234  from terminal_utils import Pager
1235  from basf2.utils import pretty_print_description_list
1236  with Pager('Flavor Tagger function accepts the following arguments:'):
1237  pretty_print_description_list(desc_list)
def isB2BII()
Definition: b2bii.py:14