Belle II Software light-2406-ragdoll
udst.py
1#!/usr/bin/env python3
2
3
10
11"""
12This module defines the uDST (user-defined data summary table) file format.
13A uDST contains all dataobjects from mdst plus Particles and ParticleLists.
14
15.. seealso:: :ref:`mdst`
16"""
17
18import basf2
19import mdst
20import pdg
21from modularAnalysis import removeParticlesNotInLists
22from b2test_utils.datastoreprinter import DataStorePrinter, PrintObjectsModule
23
24
25def add_udst_output(
26 path, filename, particleLists=None, additionalBranches=None, dataDescription=None, mc=True,
27):
28 """
29 Save uDST (user-defined Data Summary Tables) = MDST + Particles + ParticleLists
30 The charge-conjugate lists of those given in particleLists are also stored.
31 Additional Store Arrays and Relations to be stored can be specified via additionalBranches
32 list argument.
33
34 See also: `mdst.add_mdst_output`
35
36 Note:
37 that this does not reduce the amount of Particle objects saved,
38 see `udst.add_skimmed_udst_output` for a function that does.
39
40 :param basf2.Path path: Path to add the output modules to.
41 :param str filename: Name of the output file.
42 :param list(str) particleLists: Names of the particle lists to write out.
43 :param list(str) additionalBranches: datastore arrays/objects to write to the output
44 file in addition to mdst and particle information
45 :param dict dataDescription: Additional data descriptions to add to the output
46 file. For example {"mcEventType":"mixed"}
47 :param bool mc: Save Monte Carlo quantities? (MCParticles and corresponding relations)
48 """
49
50 if particleLists is None:
51 particleLists = []
52 if additionalBranches is None:
53 additionalBranches = []
54
55 # also add anti-particle lists
56 plSet = set(particleLists)
57 for List in particleLists:
58 name, label = List.split(":")
59 plSet.add(pdg.conjugate(name) + ":" + label)
60
61 # define the dataobjects to be included
62 partBranches = (
63 [
64 "Particles",
65 "ParticlesToMCParticles",
66 "ParticlesToPIDLikelihoods",
67 "ParticleExtraInfoMap",
68 "EventExtraInfo",
69 ]
70 + additionalBranches
71 + list(plSet)
72 )
73
74 # set dataDescription: dictionary is mutable and thus not a good
75 # default argument.
76 if dataDescription is None:
77 dataDescription = {}
78
79 dataDescription.update(dataLevel="udst")
80
82 path,
83 mc=mc,
84 filename=filename,
85 additionalBranches=partBranches,
86 dataDescription=dataDescription,
87 )
88
89
90def add_skimmed_udst_output(
91 path,
92 skimDecayMode,
93 skimParticleLists=None,
94 outputParticleLists=None,
95 additionalBranches=None,
96 outputFile=None,
97 dataDescription=None,
98 mc=True,
99):
100 """
101 Create a new path for events that contain a non-empty particle list specified via skimParticleLists.
102 Write the accepted events as a udst file, saving only particles from skimParticleLists
103 and from outputParticleLists.
104 Additional Store Arrays and Relations to be stored can be specified via additionalBranches
105 list argument.
106
107 :param basf2.Path path: Path to add the skim output to.
108 :param str skimDecayMode: Name of the skim. If no outputFile is given this is
109 also the name of the output filename. This name will be added to the
110 FileMetaData as an extra data description "skimDecayMode"
111 :param list(str) skimParticleLists: Names of the particle lists to skim for.
112 An event will be accepted if at least one of the particle lists is not empty
113 :param list(str) outputParticleLists: Names of the particle lists to store in
114 the output in addition to the ones in skimParticleLists
115 :param list(str) additionalBranches: datastore arrays/objects to write to the output
116 file in addition to mdst and particle information
117 :param str outputFile: Name of the output file if different from the skim name
118 :param dict dataDescription: Additional data descriptions to add to the output
119 file. For example {"mcEventType":"mixed"}
120 :param bool mc: Save Monte Carlo quantities? (MCParticles and corresponding relations)
121 """
122
123 if skimParticleLists is None:
124 skimParticleLists = []
125 if outputParticleLists is None:
126 outputParticleLists = []
127 if additionalBranches is None:
128 additionalBranches = []
129 # if no outputfile is specified, set it to the skim name
130 if outputFile is None:
131 outputFile = skimDecayMode
132
133 # make sure the output filename has the correct extension
134 if not outputFile.endswith(".udst.root"):
135 outputFile += ".udst.root"
136
137 skimfilter = basf2.register_module("SkimFilter")
138 skimfilter.set_name("SkimFilter_" + skimDecayMode)
139 skimfilter.param("particleLists", skimParticleLists)
140 path.add_module(skimfilter)
141 filter_path = basf2.create_path()
142 skimfilter.if_value("=1", filter_path, basf2.AfterConditionPath.CONTINUE)
143
144 # add_independent_path() is rather expensive, only do this for skimmed events
145 skim_path = basf2.create_path()
146 saveParticleLists = skimParticleLists + outputParticleLists
147 removeParticlesNotInLists(saveParticleLists, path=skim_path)
148
149 # set dataDescription: dictionary is mutable and thus not a good
150 # default argument.
151 if dataDescription is None:
152 dataDescription = {}
153
154 dataDescription.setdefault("skimDecayMode", skimDecayMode)
155 add_udst_output(
156 skim_path,
157 outputFile,
158 saveParticleLists,
159 additionalBranches,
160 dataDescription=dataDescription,
161 mc=mc,
162 )
163 filter_path.add_independent_path(skim_path, "skim_" + skimDecayMode)
164
165
166def add_udst_dump(path, print_untested=False):
167 """
168 Add a PrintObjectsModule to a path for printing the uDST contents to
169 test and maintain backwards compatibility.
170
171 See also: mdst.add_mdst_dump
172
173 Arguments:
174 path (basf2.Path): Path to add modules to
175 print_untested (bool): If True print the names of all methods which
176 are not explicitly printed to make sure we don't miss addition of new members
177 """
178
179 eparticlesourceobjects = [i for i in range(7)]
180
181 # Define a list of all the udst dataobjects we want to print out
182 # and all the members we want to check
183 protected_udst_dataobjects = [
184 DataStorePrinter(
185 "Particle",
186 [
187 "getPDGCode",
188 "getPDGMass",
189 "getCharge",
190 "getFlavorType",
191 "getParticleSource",
192 "getMdstArrayIndex",
193 "getProperty",
194 "getMass",
195 "getEnergy",
196 "get4Vector",
197 "getMomentum",
198 "getMomentumMagnitude",
199 "getP",
200 "getPx",
201 "getPy",
202 "getPz",
203 "getVertex",
204 "getX",
205 "getY",
206 "getZ",
207 "getPValue",
208 "getMomentumVertexErrorMatrix",
209 "getMomentumErrorMatrix",
210 "getVertexErrorMatrix",
211 "getCosHelicity",
212 "getAcoplanarity",
213 "getMdstSource",
214 "getNDaughters",
215 "getDaughterIndices",
216 "getDaughterProperties",
217 "getDaughters",
218 "getFinalStateDaughters",
219 "getTrack",
220 "getTrackFitResult",
221 "getPDGCodeUsedForFit",
222 "wasExactFitHypothesisUsed",
223 "getV0",
224 "getPIDLikelihood",
225 "getECLCluster",
226 "getECLClusterEnergy",
227 "getECLClusterEHypothesisBit",
228 "getKLMCluster",
229 "getMCParticle",
230 "getName",
231 "getInfoHTML",
232 "getExtraInfoNames",
233 "getExtraInfoMap",
234 "getExtraInfoSize",
235 "getEffectiveMomentumScale",
236 "print",
237 ],
238 {
239 "getRelationsWith": [
240 "MCParticles",
241 "TrackFitResults",
242 "PIDLikelihoods"
243 ],
244 "getMdstArrayIndices": eparticlesourceobjects,
245 "addExtraInfo": [["foo", 1337]],
246 "hasExtraInfo": ["foo", "bar"],
247 "getExtraInfo": ["foo"],
248 "getDaughter": [0, 1],
249 "getCosHelicityDaughter": [0],
250 "getParticleFromGeneralizedIndexString": ["0:0", "0:1"],
251 "updateMass": [13],
252 },
253 ),
254 DataStorePrinter(
255 "ParticleExtraInfoMap",
256 ["getNMaps"],
257 {
258 "getMap": [0],
259 "getIndex": [[0, "foo"], [0, "bar"]],
260 "getMapForNewVar": [["foo"], ["bar"]],
261 },
262 array=False,
263 ),
264 DataStorePrinter(
265 "EventExtraInfo",
266 [
267 "getNames",
268 "getInfoHTML",
269 ],
270 {
271 "addExtraInfo": [["foo", 1337]],
272 "hasExtraInfo": ["foo"],
273 "getExtraInfo": ["foo"],
274 },
275 array=False,
276 ),
277 ]
278 path.add_module(PrintObjectsModule(protected_udst_dataobjects, print_untested))
279
def add_mdst_output(path, mc=True, filename='mdst.root', additionalBranches=[], dataDescription=None)
Definition: mdst.py:38
def conjugate(name)
Definition: pdg.py:111