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