Belle II Software development
steering_training_data.py
1#!/usr/bin/env python3
2
3
10
11import os
12
13import basf2 as b2
14import modularAnalysis as ma
15from dft.DeepFlavorTagger import DeepFlavorTagger
16from vertex import treeFit, kFit
17
18
19def 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
50
51def test_expert(working_dir, file_names, identifier, output_variable='networkOutput', environmentType='default',
52 max_events=0):
53 main = b2.create_path()
54
55 ma.inputMdstList(file_names, environmentType=environmentType, path=main)
56
57 ma.fillParticleListFromMC('nu_tau:MC', '', path=main)
58 ma.reconstructMCDecay('B0:sig -> nu_tau:MC anti-nu_tau:MC', '', writeOut=True, path=main)
59
60 ma.buildRestOfEvent('B0:sig', path=main)
61
62 DeepFlavorTagger('B0:sig', 'expert', working_dir, identifier, path=main)
63
64 # define output variable
65 output_variable_name = ''.join('extraInfo(', output_variable, ')')
66
67 ma.variablesToNtuple('B0:sig', ['extraInfo(qrCombined)', output_variable_name],
68 filename=os.path.join(working_dir, identifier + '_test_output.root'),
69 path=main)
70
71 main.add_module('ProgressBar')
72
73 b2.process(main, max_events)
74
75
76def test_expert_jpsi(working_dir, file_names, prefix, environmentType='default', max_events=0):
77 main = b2.create_path()
78
79 ma.inputMdstList(file_names, environmentType=environmentType, path=main)
80
81 ma.fillParticleList('pi+:highPID', 'pionID >= .1', path=main)
82 ma.fillParticleList('mu+:highPID', 'muonID >= .1', path=main)
83
84 # reconstruct Ks -> pi+ pi- decay
85 # keep only candidates with dM<0.25
86 ma.reconstructDecay('K_S0:pipi -> pi+:highPID pi-:highPID', '.25 <= M <= .75', path=main)
87
88 # fit K_S0 Vertex
89 treeFit('K_S0:pipi', 0., path=main)
90
91 # reconstruct J/psi -> mu+ mu- decay and fit vertex
92 ma.reconstructDecay('J/psi:mumu -> mu+:highPID mu-:highPID', '3.0 <= M <= 3.2 ', path=main)
93
94 # applyCuts('J/psi:mumu', '3.07 < M < 3.11', path=main)
95 treeFit('J/psi:mumu', 0., massConstraint=['J/psi'], path=main)
96
97 # reconstruct B0 -> J/psi Ks decay
98 ma.reconstructDecay('B0:jpsiks -> J/psi:mumu K_S0:pipi', '5.2 <= M <= 5.4', path=main)
99
100 # Fit the B0 Vertex
101 kFit('B0:jpsiks', 0., 'vertex', decay_string='B0 -> [J/psi -> ^mu+ ^mu-] K_S0', path=main)
102
103 # perform MC matching (MC truth association). Always before TagV
104 ma.matchMCTruth('B0:jpsiks', path=main)
105
106 # build the rest of the event associated to the B0
107 ma.buildRestOfEvent('B0:jpsiks', path=main)
108 ma.applyCuts('B0:jpsiks', 'isSignal > 0.5', path=main)
109
110 DeepFlavorTagger('B0:jpsiks', 'Expert', working_dir, prefix, transform_to_probability=True, path=main)
111 ma.variablesToNtuple('B0:jpsiks', ['extraInfo(qrCombined)', 'extraInfo(qrMC)', 'extraInfo(B0Probability)',
112 'extraInfo(BOProbabilityMC)'],
113 filename=os.path.join(working_dir, 'test_output.root'), path=main)
114
115 main.add_module('ProgressBar')
116
117 b2.process(main, max_events)