Belle II Software light-2406-ragdoll
create_training_files.py
1#!/usr/bin/env python
2
3
10
11
12from pathlib import Path
13import argparse
14
15import basf2 as b2
16import modularAnalysis as ma
17
18from ROOT import Belle2
19
20import stdPhotons
21from variables import variables as vm
22
23from grafei import lcaSaver
24
25
26def _get_args():
27 parser = argparse.ArgumentParser(
28 formatter_class=argparse.ArgumentDefaultsHelpFormatter,
29 description='Reconstruct most likely particles and save their features and LCA matrix'
30 )
31 parser.add_argument('-t', '--type', choices=['B0', 'B+', 'Ups'], required=True,
32 help='Training target', metavar='type',
33 dest='type')
34 return parser.parse_args()
35
36
37if __name__ == '__main__':
38 b2.conditions.prepend_globaltag(ma.getAnalysisGlobaltag())
39
40 args = _get_args()
41
42 input_file = Path(Belle2.Environment.Instance().getInputFilesOverride()[0])
43
44 path = b2.create_path()
45 ma.inputMdst(str(input_file), path=path)
46
47 mc_particle_name = {
48 "B0": "B0:MC",
49 "B+": "B+:MC",
50 "Ups": "Upsilon(4S)"
51 }
52
53 # ###### BUILD MC B/Ups FOR LCA/TAGGING ######
54 ma.fillParticleListFromMC(mc_particle_name[args.type], '', path=path)
55
56 # These priors were obtained by counting truth-matched tracks in BB mixed MC
57 # It could be modified by the user if needed
58 priors = [0.068, 0.050, 0.7326, 0.1315, 0.0183, 0.00006]
59
60 charged_cuts = [f"pidIsMostLikely({','.join(str(p) for p in priors)})>0", 'nCDCHits>20',
61 'thetaInCDCAcceptance', 'abs(dz)<1.0',
62 'dr<0.5', 'p<5', 'pt>0.2']
63
64 photon_cuts = ['beamBackgroundSuppression>0.4', 'fakePhotonSuppression>0.3',
65 'abs(clusterTiming)<100', 'abs(formula(clusterTiming/clusterErrorTiming))<2.0',
66 '[[clusterReg==1 and E>0.09] or [clusterReg==2 and E>0.09] or [clusterReg==3 and E>0.14]]']
67
68 charged_lists = [f'{c}:final' for c in ['p+', 'e+', 'pi+', 'mu+', 'K+']]
69
70 ma.fillParticleLists(
71 [(c, ' and '.join(charged_cuts)) for c in charged_lists],
72 writeOut=True,
73 path=path,
74 )
75
77 listtype='tight',
78 path=path,
79 )
80 ma.getBeamBackgroundProbability("gamma:all", "MC15ri", path=path)
81 ma.getFakePhotonProbability("gamma:all", "MC15ri", path=path)
82 ma.cutAndCopyList(
83 'gamma:final',
84 'gamma:tight',
85 ' and '.join(photon_cuts),
86 writeOut=True,
87 path=path,
88 )
89
90 # Add requirements on total number of photons and charged in event
91 vm.addAlias('n_gamma_in_evt', 'nParticlesInList(gamma:final)')
92 vm.addAlias('n_p_in_evt', 'nParticlesInList(p+:final)')
93 vm.addAlias('n_e_in_evt', 'nParticlesInList(e+:final)')
94 vm.addAlias('n_mu_in_evt', 'nParticlesInList(mu+:final)')
95 vm.addAlias('n_pi_in_evt', 'nParticlesInList(pi+:final)')
96 vm.addAlias('n_K_in_evt', 'nParticlesInList(K+:final)')
97 vm.addAlias('n_charged_in_evt', 'formula(n_p_in_evt+n_e_in_evt+n_mu_in_evt+n_pi_in_evt+n_K_in_evt)')
98
99 ma.applyEventCuts('n_gamma_in_evt<20 and n_charged_in_evt<20', path=path)
100
101 # Set up variables to save to NTuple
102 save_vars = [
103 # Charged variables
104 'electronID_noSVD_noTOP',
105 'kaonID_noSVD',
106 'protonID_noSVD',
107 'muonID_noSVD',
108 'pionID_noSVD',
109 'p',
110 'pt',
111 'px',
112 'py',
113 'pz',
114 'dr',
115 'dz',
116 'x',
117 'y',
118 'z',
119 'charge',
120 # Neutral variables
121 'clusterNHits',
122 'clusterTiming',
123 'clusterE9E25',
124 ]
125
126 # Set up particle lists we'll work with
127 p_lists = charged_lists + ['gamma:final']
128 p_names = [x.split(':')[0] for x in p_lists]
129
130 # ##### TAG AND SAVE #######
131
132 # Flag each particle according to the B meson and decay it came from
133 for i, p_list in enumerate(p_lists):
134 # Match MC particles for all lists
135 ma.matchMCTruth(p_list, path=path)
136
137 lcaSaver(
138 particle_lists=p_lists,
139 features=save_vars,
140 mcparticle_list=mc_particle_name[args.type],
141 output_file=f'graFEI_train_{input_file.stem}.root',
142 path=path,
143 )
144
145 # Actually run everything
146 b2.process(path)
147 print(b2.statistics)
static Environment & Instance()
Static method to get a reference to the Environment instance.
Definition: Environment.cc:28
def stdPhotons(listtype='loose', path=None, beamBackgroundMVAWeight="", fakePhotonMVAWeight="", biasCorrectionTable="")
Definition: stdPhotons.py:19