Belle II Software light-2406-ragdoll
__init__.py
1
8
9
10__all__ = ["graFEI", "lcaSaver"]
11
12
13from grafei.modules.GraFEIModule import GraFEIModule
14from grafei.modules.LCASaverModule import LCASaverModule
15
16import modularAnalysis as ma
17import basf2 as b2
18
19
20def graFEI(
21 list_name,
22 path,
23 particle_lists=None,
24 store_mc_truth=False,
25 cfg_path=None,
26 param_file=None,
27 sig_side_lcas=None,
28 sig_side_masses=None,
29 gpu=False,
30 payload_config_name="graFEIConfigFile",
31 payload_model_name="graFEIModelFile",
32):
33 """
34 Wrapper function to add the GraFEIModule to the path and perform other actions behind the scenes.
35
36 Applies graFEI model to a (list of) particle list(s) in basf2.
37 GraFEI information is stored as eventExtraInfos.
38
39 .. note::
40 ``list_name`` should always be provided. This is the name of the particle list to be given as input to the graFEI.
41 If ``list_name`` refers to an existing particle list, it is used as input to the model.
42 If also a list of final state particle lists is provided in ``particle_lists``, these are combined to form
43 a new list called ``list_name`` (if ``list_name`` already exists an error is thrown).
44 If ``particle_list`` is provided, the mass hypotheses of final state particles are updated to match graFEI predictions.
45
46 Args:
47 list_name (str): Name of particle list given as input to the model.
48 path (basf2.Path): Module is added to this path.
49 particle_lists (list): List of particle lists. If provided, these are combined to form ``list_name``.
50 store_mc_truth (bool): Whether to store MC truth information.
51 cfg_path (str): Path to config file. If `None` the config file in the global tag is used.
52 param_file (str): Path to parameter file containing the model. If `None` the parameter file in the global tag is used.
53 sig_side_lcas (list): List containing LCAS matrix of signal-side.
54 sig_side_masses (list): List containing mass hypotheses of signal-side.
55 gpu (bool): Whether to run on a GPU.
56 payload_config_name (str): Name of config file payload. The default should be kept, except in basf2 examples.
57 payload_model_name (str): Name of model file payload. The default should be kept, except in basf2 examples.
58
59 Returns:
60 list: List of graFEI variables.
61 """
62 if particle_lists:
63 ma.combineAllParticles(particle_lists, list_name, path=path)
64 if store_mc_truth:
65 ma.matchMCTruth(list_name, path=path)
66 ma.fillParticleListFromMC("Upsilon(4S):MC", "", path=path)
67 ma.fillParticleListFromMC("B0:MC", "", path=path)
68 ma.fillParticleListFromMC("B+:MC", "", path=path)
69
70 graFEI = GraFEIModule(
71 list_name,
72 cfg_path=cfg_path,
73 param_file=param_file,
74 sig_side_lcas=sig_side_lcas,
75 sig_side_masses=sig_side_masses,
76 gpu=gpu,
77 payload_config_name=payload_config_name,
78 payload_model_name=payload_model_name,
79 )
80 path.add_module(graFEI)
81
82 graFEI_vars = [
83 "graFEI_probEdgeProd",
84 "graFEI_probEdgeMean",
85 "graFEI_probEdgeGeom",
86 "graFEI_validTree",
87 "graFEI_goodEvent",
88 "graFEI_nFSP",
89 "graFEI_nCharged_preFit",
90 "graFEI_nElectrons_preFit",
91 "graFEI_nMuons_preFit",
92 "graFEI_nPions_preFit",
93 "graFEI_nKaons_preFit",
94 "graFEI_nProtons_preFit",
95 "graFEI_nLeptons_preFit",
96 "graFEI_nPhotons_preFit",
97 "graFEI_nOthers_preFit",
98 "graFEI_nCharged_postFit",
99 "graFEI_nElectrons_postFit",
100 "graFEI_nMuons_postFit",
101 "graFEI_nPions_postFit",
102 "graFEI_nKaons_postFit",
103 "graFEI_nProtons_postFit",
104 "graFEI_nLeptons_postFit",
105 "graFEI_nPhotons_postFit",
106 "graFEI_nOthers_postFit",
107 "graFEI_nPredictedUnmatched",
108 "graFEI_nPredictedUnmatched_noPhotons",
109 ]
110 if store_mc_truth:
111 graFEI_vars.extend(
112 [
113 "graFEI_truth_perfectLCA",
114 "graFEI_truth_perfectMasses",
115 "graFEI_truth_perfectEvent",
116 "graFEI_truth_isSemileptonic",
117 "graFEI_truth_nFSP",
118 "graFEI_truth_nPhotons",
119 "graFEI_truth_nElectrons",
120 "graFEI_truth_nMuons",
121 "graFEI_truth_nPions",
122 "graFEI_truth_nKaons",
123 "graFEI_truth_nProtons",
124 "graFEI_truth_nOthers",
125 ]
126 )
127
128 ma.variablesToEventExtraInfo(
129 list_name,
130 dict((f"extraInfo({var})", var) for var in graFEI_vars),
131 path=path,
132 )
133
134 # Update mass hypotheses
135 if particle_lists:
136 charged_lists = [ls for ls in particle_lists if "gamma:" not in ls]
137 photon_lists = [ls for ls in particle_lists if "gamma:" in ls]
138
139 hypotheses = { # PDG code and graFEI class for each mass hypothesis
140 "K": (321, 4),
141 "pi": (211, 3),
142 "mu": (13, 2),
143 "e": (11, 1),
144 "p": (2212, 5),
145 }
146 for newHyp in hypotheses: # Loop over all possible final mass hypotheses
147 newPDG = hypotheses[newHyp][0]
148 newClass = hypotheses[newHyp][1]
149 final_lists = []
150 for oldList in charged_lists: # Loop over all charged particle lists
151 ma.cutAndCopyList( # Retain only particles with predicted mass hypothesis equal to `newClass`
152 f"{oldList}_to_{newHyp}_",
153 oldList,
154 f"extraInfo(graFEI_massHypothesis) == {newClass}",
155 path=path,
156 )
157 ma.updateMassHypothesis( # Update basf2 mass hypothesis
158 f"{oldList}_to_{newHyp}_", newPDG, path=path
159 )
160
161 label = oldList[oldList.find(':')+1:]
162 oldHyp = oldList[:oldList.find(':')-1]
163 final_lists.append(f"{newHyp}+:{label}_to_{newHyp}__converted_from_{oldHyp}")
164
165 ma.copyLists( # Merge all temporary lists into final `newHyp` list
166 f"{newHyp}+:graFEI",
167 final_lists,
168 writeOut=True,
169 path=path,
170 )
171
172 # Take care of photons
173 if len(photon_lists) > 0:
174 ma.cutAndCopyList(
175 "gamma:graFEI",
176 photon_lists[0],
177 "extraInfo(graFEI_massHypothesis) == 6",
178 writeOut=True,
179 path=path,
180 )
181 elif len(photon_lists) == 0:
182 b2.B2WARNING("grafei.graFEI You did not define a photon input list. Therefore you don't have any as output.")
183 if len(photon_lists) > 1:
184 b2.B2WARNING("grafei.graFEI You defined more than one photon input list. Using the first one.")
185
186 return graFEI_vars
187
188
189def lcaSaver(
190 particle_lists,
191 features,
192 mcparticle_list,
193 output_file,
194 path,
195):
196 """
197 Wrapper function to add the LCASaverModule to the path.
198
199 Save Lowest Common Ancestor matrix of each MC Particle in the given list.
200
201 Args:
202 particle_lists (list): Name of particle lists to save features of.
203 features (list): List of features to save for each particle.
204 mcparticle_list (str): Name of particle list to build LCAs from (used as root).
205 output_file (str): Path to output file to save.
206 path (basf2.Path): Module is added to this path.
207 """
208 root_saver_module = LCASaverModule(
209 particle_lists=particle_lists,
210 features=features,
211 mcparticle_list=mcparticle_list,
212 output_file=output_file,
213 )
214 path.add_module(root_saver_module)
215
216
217print(
218 r"""
219 ____ ____ _ ____ ____ _
220 | _ |__| /_\ |___ |___ |
221 |__| | \ / \ | |___ |
222
223 o
224 / \
225 / \ x-----x _ _
226 / \ |-----\ |\ /| |-----\ |0 1 3 3 5|
227 o \ |----- \ | \ / | |----- \ |1 0 3 3 5|
228 / \ \ |----- / | x | |----- / |3 3 0 1 5|
229 / \ \ |-----/ | / \ | |-----/ |3 3 1 0 5|
230 / \ \ |/ \| |5 5 5 5 0|
231 o o \ x-----x  ̄  ̄
232 / \ / \ \
233 x x x x x
234
235 Authors: Jacopo Cerasoli, Giulio Dujany, Lucas Martel, Corentin Santos. 2022 - 2024
236 Model description: https://indico.cern.ch/event/1106990/papers/4996235/files/12252-ACAT_2022_proceedings.pdf
237 Based on the work of Kahn et al: https://iopscience.iop.org/article/10.1088/2632-2153/ac8de0
238 Please consider citing both articles.
239 Code adapted from https://github.com/Helmholtz-AI-Energy/BaumBauen
240 """
241)