Belle II Software light-2406-ragdoll
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 print(b2.statistics)
50
51
52def 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
78def 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)