Belle II Software  release-08-01-10
steering_training_data.py
1 #!/usr/bin/env python3
2 
3 
10 
11 import os
12 
13 import basf2 as b2
14 import modularAnalysis as ma
15 from dft.DeepFlavorTagger import DeepFlavorTagger
16 from vertex import treeFit, kFit
17 
18 
19 def create_train_data(
20  working_dir,
21  file_names,
22  identifier,
23  variable_list,
24  environmentType='default',
25  target='qrCombined',
26  overwrite=False,
27  max_events=0,
28  mode='sampler',
29  *args,
30  **kwargs):
31  main = b2.create_path()
32 
33  if not os.path.exists(working_dir) and working_dir != '':
34  os.makedirs(working_dir)
35 
36  ma.inputMdstList(environmentType=environmentType, filelist=file_names, path=main)
37 
38  ma.fillParticleListFromMC('nu_tau:MC', '', path=main)
39  ma.reconstructMCDecay('B0:sig -> nu_tau:MC anti-nu_tau:MC', '', writeOut=True, path=main)
40 
41  ma.buildRestOfEvent('B0:sig', path=main)
42 
43  DeepFlavorTagger('B0:sig', mode, working_dir, identifier, variable_list, target=target, overwrite=overwrite,
44  path=main, *args, **kwargs)
45 
46  main.add_module('ProgressBar')
47 
48  b2.process(main, max_events)
49  print(b2.statistics)
50 
51 
52 def test_expert(working_dir, file_names, identifier, output_variable='networkOutput', environmentType='default',
53  max_events=0):
54  main = b2.create_path()
55 
56  ma.inputMdstList(file_names, environmentType=environmentType, path=main)
57 
58  ma.fillParticleListFromMC('nu_tau:MC', '', path=main)
59  ma.reconstructMCDecay('B0:sig -> nu_tau:MC anti-nu_tau:MC', '', writeOut=True, path=main)
60 
61  ma.buildRestOfEvent('B0:sig', path=main)
62 
63  DeepFlavorTagger('B0:sig', 'expert', working_dir, identifier, path=main)
64 
65  # define output variable
66  output_variable_name = ''.join('extraInfo(', output_variable, ')')
67 
68  ma.variablesToNtuple('B0:sig', ['extraInfo(qrCombined)', output_variable_name],
69  filename=os.path.join(working_dir, identifier + '_test_output.root'),
70  path=main)
71 
72  main.add_module('ProgressBar')
73 
74  b2.process(main, max_events)
75  print(b2.statistics)
76 
77 
78 def test_expert_jpsi(working_dir, file_names, prefix, environmentType='default', max_events=0):
79  main = b2.create_path()
80 
81  ma.inputMdstList(file_names, environmentType=environmentType, path=main)
82 
83  ma.fillParticleList('pi+:highPID', 'pionID >= .1', path=main)
84  ma.fillParticleList('mu+:highPID', 'muonID >= .1', path=main)
85 
86  # reconstruct Ks -> pi+ pi- decay
87  # keep only candidates with dM<0.25
88  ma.reconstructDecay('K_S0:pipi -> pi+:highPID pi-:highPID', '.25 <= M <= .75', path=main)
89 
90  # fit K_S0 Vertex
91  treeFit('K_S0:pipi', 0., path=main)
92 
93  # reconstruct J/psi -> mu+ mu- decay and fit vertex
94  ma.reconstructDecay('J/psi:mumu -> mu+:highPID mu-:highPID', '3.0 <= M <= 3.2 ', path=main)
95 
96  # applyCuts('J/psi:mumu', '3.07 < M < 3.11', path=main)
97  treeFit('J/psi:mumu', 0., massConstraint=['J/psi'], path=main)
98 
99  # reconstruct B0 -> J/psi Ks decay
100  ma.reconstructDecay('B0:jpsiks -> J/psi:mumu K_S0:pipi', '5.2 <= M <= 5.4', path=main)
101 
102  # Fit the B0 Vertex
103  kFit('B0:jpsiks', 0., 'vertex', decay_string='B0 -> [J/psi -> ^mu+ ^mu-] K_S0', path=main)
104 
105  # perform MC matching (MC truth association). Always before TagV
106  ma.matchMCTruth('B0:jpsiks', path=main)
107 
108  # build the rest of the event associated to the B0
109  ma.buildRestOfEvent('B0:jpsiks', path=main)
110  ma.applyCuts('B0:jpsiks', 'isSignal > 0.5', path=main)
111 
112  DeepFlavorTagger('B0:jpsiks', 'Expert', working_dir, prefix, transform_to_probability=True, path=main)
113  ma.variablesToNtuple('B0:jpsiks', ['extraInfo(qrCombined)', 'extraInfo(qrMC)', 'extraInfo(B0Probability)',
114  'extraInfo(BOProbabilityMC)'],
115  filename=os.path.join(working_dir, 'test_output.root'), path=main)
116 
117  main.add_module('ProgressBar')
118 
119  b2.process(main, max_events)
120  print(b2.statistics)