Belle II Software light-2406-ragdoll
trackIsolationVariables.py
1#!/usr/bin/env python3
2
3
10
11# Doxygen should skip this script
12# @cond
13
14"""
15Example script to calculate isolation variables per particle.
16
17For each particle in the input charged stable particle list,
18calculate the distance to the closest candidate in the reference list at a given detector layer surface.
19The calculation of the distance is based on the particles' track helices extrapolation.
20"""
21
22
23import argparse
24
25
26def argparser():
27 """ Argument parser
28 """
29
30 import stdCharged as stdc
31 from modularAnalysis import getAnalysisGlobaltag
32
33 parser = argparse.ArgumentParser(description=__doc__,
34 formatter_class=argparse.RawTextHelpFormatter)
35
36 parser.add_argument("--std_charged_ref",
37 type=str,
38 choices=stdc._chargednames,
39 default="pi",
40 help="The base name of the reference standard charged particle list\n"
41 "that will be considered for the distance calculation.\n"
42 "Default: %(default)s.")
43 parser.add_argument("--detectors",
44 type=str,
45 nargs="+",
46 default=["CDC", "TOP", "ARICH", "ECL", "KLM"],
47 choices=["CDC", "TOP", "ARICH", "ECL", "KLM"],
48 help="List of detectors at whose entry surface the isolation variables will be calculated.\n"
49 "Pass a space-separated list of names.\n"
50 "Default: %(default)s.")
51 parser.add_argument("--use_pid_det_weights",
52 action="store_true",
53 default=False,
54 help="Include the PID detector weights (taken from the CDB) in the isolation score calculation.\n"
55 "Default: %(default)s.")
56 parser.add_argument("--global_tag_append",
57 type=str,
58 nargs="+",
59 default=[getAnalysisGlobaltag()],
60 help="List of names of conditions DB global tag(s) to append on top of GT replay.\n"
61 "NB: these GTs will have lowest priority over GT replay.\n"
62 "The order of the sequence passed determines the priority of the GTs, w/ the highest coming first.\n"
63 "Pass a space-separated list of names.\n"
64 "Default: %(default)s.")
65 parser.add_argument("--global_tag_prepend",
66 type=str,
67 nargs="+",
68 default=None,
69 help="List of names of conditions DB global tag(s) to prepend to GT replay.\n"
70 "NB: these GTs will have highest priority over GT replay.\n"
71 "The order of the sequence passed determines the priority of the GTs, w/ the highest coming first.\n"
72 "Pass a space-separated list of names.")
73 parser.add_argument("-d", "--debug",
74 action="store",
75 default=0,
76 type=int,
77 choices=list(range(11, 20)),
78 help="Run the TrackIsoCalculator module in debug mode. Pass the desired DEBUG level integer.")
79
80 return parser
81
82
83if __name__ == "__main__":
84
85 # Argparse options.
86 # NB: Command line arguments are parsed before importing basf2, to avoid PyROOT hijacking them
87 # in case of overlapping option names.
88 args = argparser().parse_args()
89
90 import basf2 as b2
91 import modularAnalysis as ma
92 import stdV0s as stdv0
93 from variables import variables as vm
94 import variables.utils as vu
95 import variables.collections as vc
96 import pdg
97 from ROOT import Belle2
98 Const = Belle2.Const
99
100 for tag in args.global_tag_append:
101 b2.conditions.append_globaltag(tag)
102 print(f"Appending GTs:\n{args.global_tag_append}")
103
104 if args.global_tag_prepend:
105 for tag in reversed(args.global_tag_prepend):
106 b2.conditions.prepend_globaltag(tag)
107 print(f"Prepending GTs:\n{args.global_tag_prepend}")
108
109 # Create path. Register necessary modules to this path.
110 path = b2.create_path()
111
112 # Add input data and ParticleLoader modules to the path.
113 ma.inputMdst(filename=b2.find_file("mdst14.root", "validation"), path=path)
114
115 # -------------------------------------------------------------------
116
117 # ---------
118 # EXAMPLE 1
119 # ---------
120
121 # Fill a particle list of muons, with some quality selection.
122 base_trk_selection = "[dr < 3] and [abs(dz) < 5] and [thetaInCDCAcceptance] and [pt > 0.1]"
123 ma.fillParticleList("mu+:muons", f"{base_trk_selection} and [muonID > 0.5]", path=path)
124
125 # Reconstruct a J/psi decay.
126 jpsimumu = "J/psi:mumu -> mu+:muons mu-:muons"
127 jpsi_cuts = [
128 "[2.8 < M < 3.3]",
129 "[daughterSumOf(charge) == 0]",
130 ]
131 jpsi_cut = " and ".join(jpsi_cuts)
132
133 ma.reconstructDecay(jpsimumu, jpsi_cut, path=path)
134
135 # Fill a reference charged stable particle list to calculate the distances against.
136 # Generally, this list should have a very loose selection applied (if none at all).
137 ref = f"{args.std_charged_ref}+:ref"
138 ref_pdg = pdg.from_name(f"{args.std_charged_ref}+")
139 ma.fillParticleList(ref, f"{base_trk_selection}", path=path)
140
141 # Calculate the track isolation variables
142 # on the muon daughters in the decay.
143 # Use the default setting where the module takes the mass hypothesis w/ highest probability for the track extrapolation.
144 # The helper returns a dictionary w/ the list of metavariables for each of the reference particle lists (using PDG as key).
145 # In this case, we pass explicitly a particle list of our choice.
146 trackiso_vars_mu = ma.calculateTrackIsolation("J/psi:mumu -> ^mu+ ^mu-",
147 path,
148 *args.detectors,
149 vars_for_nearest_part=vc.mc_variables,
150 # Calculate also the chosen variables for the nearest particle at each layer.
151 reference_list_name=ref,
152 # Include/exclude the PID detector weights in the score calculation.
153 exclude_pid_det_weights=not args.use_pid_det_weights)
154
155 # Variables and aliases for the J/psi candidates.
156 variables_jpsi = vc.kinematics + ["daughterDiffOfPhi(0, 1)"]
157 variables_jpsi += vu.create_aliases(variables_jpsi, "useCMSFrame({variable})", "CMS")
158 variables_jpsi += vc.inv_mass
159 aliases_jpsi = vu.create_aliases_for_selected(variables_jpsi,
160 "^J/psi:mumu -> mu+ mu-",
161 prefix=["jpsi"])
162
163 # Variables and aliases for the J/psi daughters.
164 # Since we passed explicitly a reference list, we use its associated PDG code to get the list of metavariables.
165 variables_mu = vc.kinematics + ["theta", "phi", "clusterE", "nCDCHits"] + trackiso_vars_mu[ref_pdg]
166 aliases_mu = vu.create_aliases_for_selected(variables_mu,
167 "J/psi:mumu -> ^mu+ ^mu-",
168 use_names=True)
169
170 # Variables and aliases for the event.
171 vm.addAlias("nReferenceTracks", f"nCleanedTracks({base_trk_selection})")
172 aliases_event = ["nReferenceTracks"]
173
174 # Saving variables to ntuple
175 ma.variablesToNtuple(decayString="J/psi:mumu",
176 variables=aliases_event+aliases_jpsi+aliases_mu,
177 treename="jpsimumu",
178 filename="TrackIsolationVariables.root",
179 path=path)
180
181 # ---------
182 # EXAMPLE 2
183 # ---------
184
185 # Reconstruct standard Lambda0 -> p+ pi- decays.
186 stdv0.stdLambdas(path=path)
187
188 # Calculate the track isolation variables
189 # on the proton and pion daughters in the decay.
190 #
191 # In this configuration, the mass hypothesis for the extrapolation is the one matching each particle's PDG.
192 #
193 # Note that no reference list is passed: it will use by default the `:all` ParticleList of the same type
194 # of the selected particle(s) in the decay string.
195 trackiso_vars_p_pi = ma.calculateTrackIsolation("Lambda0:merged -> ^p+ ^pi-",
196 path,
197 *args.detectors,
198 vars_for_nearest_part=vc.mc_variables,
199 # Calculate also the chosen variables for the nearest particle at each layer.
200 highest_prob_mass_for_ext=False,
201 exclude_pid_det_weights=not args.use_pid_det_weights)
202
203 # Variables and aliases for the Lambda0 candidates.
204 variables_lambda0 = vc.kinematics + ["daughterDiffOfPhi(0, 1)"]
205 variables_lambda0 += vu.create_aliases(variables_lambda0, "useCMSFrame({variable})", "CMS")
206 variables_lambda0 += vc.inv_mass
207 aliases_lambda0 = vu.create_aliases_for_selected(variables_lambda0,
208 "^Lambda0:merged -> p+ pi-",
209 prefix=["lambda0"])
210
211 # Variables and aliases for the Lambda0 daughters.
212 # - Protons
213 variables_p = vc.kinematics + \
214 ["theta", "phi", "clusterE", "nCDCHits"] + \
215 trackiso_vars_p_pi[Const.proton.getPDGCode()] # Use the proton PDG to get the associated list of metavariables.
216 aliases_p = vu.create_aliases_for_selected(variables_p,
217 "Lambda0:merged -> ^p+ pi-",
218 use_names=True)
219 # - Pions
220 variables_pi = vc.kinematics + \
221 ["theta", "phi", "clusterE", "nCDCHits"] + \
222 trackiso_vars_p_pi[Const.pion.getPDGCode()] # Use the pion PDG to get the associated list of metavariables.
223 aliases_pi = vu.create_aliases_for_selected(variables_pi,
224 "Lambda0:merged -> p+ ^pi-",
225 use_names=True)
226
227 # Saving variables to ntuple
228 ma.variablesToNtuple(decayString="Lambda0:merged",
229 variables=aliases_lambda0+aliases_p+aliases_pi,
230 treename="lambda0ppi",
231 filename="TrackIsolationVariables.root",
232 path=path)
233
234 # -------------------------------------------------------------------
235
236 vm.printAliases()
237
238 # Optionally activate debug mode for the TrackIsoCalculator module(s).
239 if args.debug:
240 for m in path.modules():
241 if "TrackIsoCalculator" in m.name():
242 m.set_log_level(b2.LogLevel.DEBUG)
243 m.set_debug_level(args.debug)
244
245 path.add_module("Progress")
246
247 # Process the data.
248 b2.process(path)
249
250 print(b2.statistics)
251
252# @endcond
This class provides a set of constants for the framework.
Definition: Const.h:34
def from_name(name)
Definition: pdg.py:63