Belle II Software light-2406-ragdoll
graFEI_ntupliser_Upsilon.py
1#!/usr/bin/env python3
2
3
10
11
12import basf2 as b2
13import modularAnalysis as ma
14from variables import variables as vm
15import stdPhotons
16
17# Necessary to run argparse
18from ROOT import PyConfig
19PyConfig.IgnoreCommandLineOptions = True # noqa
20
21import random
22import argparse
23
24from grafei import graFEI
25
26
27# Random seeds
28b2.set_random_seed(42)
29random.seed(42)
30
31
32def _parse_args():
33 parser = argparse.ArgumentParser()
34 parser.add_argument(
35 "-c", "--config", type=str, default=None, help="graFEI config file"
36 )
37 parser.add_argument(
38 "-w", "--weight", type=str, default=None, help="graFEI weight file"
39 )
40 parser.add_argument(
41 "-g",
42 "--globaltag",
43 type=str,
44 default="analysis_tools_light-2403-persian",
45 help="Globaltag containing graFEI model",
46 )
47 parser.add_argument(
48 "-l",
49 "--lcas",
50 type=list,
51 default=[[0]],
52 help="Choose LCAS matrix of signal side",
53 )
54 parser.add_argument(
55 "-m",
56 "--masses",
57 type=list,
58 default=["k"],
59 help="Choose mass hypotheses of signal side",
60 )
61 parser.add_argument(
62 "-n",
63 "--no_mc_truth",
64 action="store_true",
65 help="Choose not to store MC-truth information",
66 )
67 return parser.parse_args()
68
69
70if __name__ == "__main__":
71 args = _parse_args()
72
73 store_mc_truth = not args.no_mc_truth
74
75 path = b2.create_path()
76 ma.inputMdst(filename=b2.find_file('mdst14.root', 'validation', False), path=path)
77
78 b2.conditions.prepend_globaltag(ma.getAnalysisGlobaltag())
79 if args.globaltag:
80 b2.conditions.prepend_globaltag(args.globaltag)
81
82
87
88 # 1) Cuts on charged tracks and photons
89 # These priors were obtained by counting truth-matched tracks in BB mixed MC
90 # It could be modified by the user if needed
91 priors = [0.068, 0.050, 0.7326, 0.1315, 0.0183, 0.00006]
92
93 charged_cuts = [
94 f"pidIsMostLikely({','.join(str(p) for p in priors)})>0",
95 "nCDCHits>20",
96 "thetaInCDCAcceptance",
97 "abs(dz)<1.0",
98 "dr<0.5",
99 "p<5",
100 "pt>0.2",
101 ]
102
103 photon_cuts = [
104 "beamBackgroundSuppression>0.4",
105 "fakePhotonSuppression>0.3",
106 "abs(clusterTiming)<100",
107 "abs(formula(clusterTiming/clusterErrorTiming))<2.0",
108 "[[clusterReg==1 and E>0.09] or [clusterReg==2 and E>0.09] or [clusterReg==3 and E>0.14]]",
109 ]
110
111 evt_cut = "n_gamma_in_evt<20 and n_charged_in_evt<20"
112
113 # Fill particle list with optimized cuts
114 charged_lists = [f"{c}:final" for c in ["p+", "e+", "pi+", "mu+", "K+"]]
115
116 ma.fillParticleLists(
117 [(c, " and ".join(charged_cuts)) for c in charged_lists],
118 writeOut=True,
119 path=path,
120 )
121
123 listtype="tight",
124 beamBackgroundMVAWeight="MC15ri",
125 fakePhotonMVAWeight="MC15ri",
126 path=path,
127 )
128
129 ma.cutAndCopyList(
130 "gamma:final",
131 "gamma:tight",
132 " and ".join(photon_cuts),
133 writeOut=True,
134 path=path,
135 )
136
137 # Add requirements on total number of photons and charged in event
138 vm.addAlias("n_gamma_in_evt", "nParticlesInList(gamma:final)")
139 vm.addAlias("n_p_in_evt", "nParticlesInList(p+:final)")
140 vm.addAlias("n_e_in_evt", "nParticlesInList(e+:final)")
141 vm.addAlias("n_mu_in_evt", "nParticlesInList(mu+:final)")
142 vm.addAlias("n_pi_in_evt", "nParticlesInList(pi+:final)")
143 vm.addAlias("n_K_in_evt", "nParticlesInList(K+:final)")
144 vm.addAlias(
145 "n_charged_in_evt",
146 "formula(n_p_in_evt+n_e_in_evt+n_mu_in_evt+n_pi_in_evt+n_K_in_evt)",
147 )
148
149 ma.applyEventCuts(evt_cut, path=path)
150
151 particle_lists = charged_lists + ["gamma:final"]
152 particle_types = [x.split(":")[0] for x in particle_lists]
153 charged_types = [x.split(":")[0] for x in charged_lists]
154
155 graFEI_vars = graFEI(
156 "Upsilon(4S):final",
157 particle_lists=particle_lists,
158 store_mc_truth=store_mc_truth,
159 cfg_path=args.config,
160 param_file=args.weight,
161 sig_side_lcas=args.lcas,
162 sig_side_masses=args.masses,
163 payload_config_name="graFEIConfigFile_Upsreco_example", # If you use default payload name just remove this argument
164 payload_model_name="graFEIModelFile_Upsreco_example", # If you use default payload name just remove this argument
165 path=path,
166 )
167
168 # Define signal-side B
169 ma.reconstructDecay(
170 "B+:Bsgn -> K+:graFEI",
171 "daughter(0, extraInfo(graFEI_sigSide)) == 1",
172 path=path,
173 )
174
175 # Define tag-side B
176 for part in particle_types:
177 ma.cutAndCopyList(
178 f"{part}:Btag",
179 f"{part}:graFEI",
180 cut="extraInfo(graFEI_sigSide) == 0",
181 writeOut=True,
182 path=path,
183 )
184 ma.combineAllParticles([f"{part}:Btag" for part in particle_types], "B+:Btag", path=path)
185
186 # 3) Keep only "good events" i.e. validTree + signal side matches chosen LCAS and mass hypotheses
187 ma.reconstructDecay(
188 "Upsilon(4S):neutral -> B+:Bsgn B-:Btag", "", path=path
189 )
190 ma.reconstructDecay(
191 "Upsilon(4S):charged -> B+:Bsgn B+:Btag", "", allowChargeViolation=True, path=path
192 )
193
194 ma.copyLists(
195 "Upsilon(4S):graFEI",
196 [
197 "Upsilon(4S):neutral",
198 "Upsilon(4S):charged",
199 ],
200 path=path,
201 )
202
203 # Reject events with no signal candidates
204 skimfilter = b2.register_module("SkimFilter")
205 skimfilter.param("particleLists", ["Upsilon(4S):graFEI"])
206 empty_path = b2.create_path()
207 skimfilter.if_value("=0", empty_path, b2.AfterConditionPath.END)
208 path.add_module(skimfilter)
209
210 if store_mc_truth:
211 ma.matchMCTruth("B+:Bsgn", path=path)
212 ma.matchMCTruth("B-:Btag", path=path)
213 ma.matchMCTruth("Upsilon(4S):graFEI", path=path)
214
215
218
219 # Variables
220 momentum_vars = [
221 "p",
222 "px",
223 "py",
224 "pz",
225 "pt",
226 "E",
227 ]
228 default_vars = [
229 "PDG",
230 "M",
231 "Mbc",
232 "deltaE",
233 ]
234 default_vars += momentum_vars
235
236 tm_vars = [
237 "mcErrors",
238 "genMotherPDG",
239 "mcPDG",
240 ] if store_mc_truth else []
241
242 default_vars += tm_vars
243
244 # Aliases
245 for var in default_vars:
246 vm.addAlias(f"Bsgn_{var}", f"daughter(0, {var})")
247 vm.addAlias(f"Bsgn_d0_{var}", f"daughter(0, daughter(0, {var}))")
248 vm.addAlias(f"Btag_{var}", f"daughter(1, {var})")
249 vm.addAlias(f"Ups_{var}", var)
250 for var in momentum_vars:
251 vm.addAlias(
252 f"Bsgn_d0_CMS_{var}", f"useCMSFrame(daughter(0, daughter(0, {var})))"
253 )
254 vm.addAlias(f"Btag_CMS_{var}", f"useCMSFrame(daughter(1, {var}))")
255 for var in graFEI_vars:
256 vm.addAlias(var, f"eventExtraInfo({var})")
257
258 all_vars = (
259 graFEI_vars
260 + [f"Bsgn_{var}" for var in tm_vars]
261 + [f"Bsgn_d0_{var}" for var in default_vars]
262 + [f"Bsgn_d0_CMS_{var}" for var in momentum_vars]
263 + [f"Btag_{var}" for var in default_vars]
264 + [f"Btag_CMS_{var}" for var in momentum_vars]
265 + [f"Ups_{var}" for var in default_vars]
266 )
267
268 ma.variablesToNtuple(
269 "Upsilon(4S):graFEI",
270 sorted(all_vars),
271 filename="graFEI_UpsReco_example.root",
272 treename="tree",
273 path=path,
274 )
275
276 # Process
277 b2.process(path)
278
279 # print out the summary
280 print(b2.statistics)
def stdPhotons(listtype='loose', path=None, beamBackgroundMVAWeight="", fakePhotonMVAWeight="", biasCorrectionTable="")
Definition: stdPhotons.py:19