Belle II Software light-2406-ragdoll
mdst.py
1#!/usr/bin/env python3
2
3
10
11from b2test_utils.datastoreprinter import DataStorePrinter, PrintObjectsModule
12
13
14MDST_OBJECTS = (
15 'ECLClusters',
16 'ECLClustersToTracksNamedBremsstrahlung',
17 'EventLevelClusteringInfo',
18 'EventLevelTrackingInfo',
19 'EventLevelTriggerTimeInfo',
20 'KLMClusters',
21 'Kinks',
22 'KlIds',
23 'PIDLikelihoods',
24 'SoftwareTriggerResult',
25 'TrackFitResults',
26 'Tracks',
27 'TRGSummary',
28 'V0s',
29)
30
31
32def add_mdst_output(
33 path,
34 mc=True,
35 filename='mdst.root',
36 additionalBranches=[],
37 dataDescription=None,
38):
39 """
40 Add the mDST output module to a path.
41 This function defines the mDST data format.
42
43 Arguments:
44 path (basf2.Path): Path to add module to
45 mc (bool): Save Monte Carlo quantities? (MCParticles and corresponding relations)
46 filename (str): Output file name.
47 additionalBranches (list): Additional objects/arrays of event durability to save
48 dataDescription (dict or None): Additional key->value pairs to be added as data description
49 fields to the output FileMetaData
50 """
51 branches = list(MDST_OBJECTS)
52 persistentBranches = ['FileMetaData']
53 if mc:
54 branches += ['MCParticles']
55 persistentBranches += ['BackgroundInfo']
56 branches += additionalBranches
57 # set dataDescription correctly
58 if dataDescription is None:
59 dataDescription = {}
60 # set dataLevel to mdst if it's not already set to something else (which
61 # might happen for udst output since that calls this function)
62 dataDescription.setdefault("dataLevel", "mdst")
63 return path.add_module("RootOutput", outputFileName=filename, branchNames=branches,
64 branchNamesPersistent=persistentBranches, additionalDataDescription=dataDescription)
65
66
67def add_mdst_dump(path, print_untested=False):
68 """
69 Add a PrintObjectsModule to a path for printing the mDST content.
70
71 Arguments:
72 path (basf2.Path): Path to add module to
73 print_untested (bool): If True print the names of all methods which are not
74 explicitly printed to make sure we don't miss addition of new members """
75
76 # Always avoid the top-level 'import ROOT'.
77 from ROOT.Belle2 import Const # noqa
78
79 # prepare a list of PID detector sets and charged stable particles
80 pid_detectors = []
81 it = Const.PIDDetectors.c_set.begin()
82 while True:
83 pid_detectors.append(Const.PIDDetectorSet(it.getDetector()))
84 it.increment()
85 if it == Const.PIDDetectors.c_set.end():
86 break
87 charged_stables = [Const.ChargedStable(Const.chargedStableSet.at(index)) for index in range(Const.chargedStableSet.size())]
88
89 # Now we define a list of all the mdst_dataobjects we want to print out and all
90 # the members we want to check
91 mdst_dataobjects = [
92 DataStorePrinter("EventMetaData", [
93 "getErrorFlag", "getEvent", "getRun", "getSubrun", "getExperiment",
94 "getProduction", "getTime", "getParentLfn", "getGeneratedWeight",
95 "isEndOfRun"
96 ], array=False),
97
98 DataStorePrinter("Track", [
99 "getNumberOfFittedHypotheses", "getQualityIndicator", "isFlippedAndRefitted",
100 "getTrackTime", "wasRefined"
101 ], {
102 "getTrackFitResult": charged_stables,
103 "getTrackFitResultWithClosestMass": charged_stables,
104 "getRelationsWith": ["ECLClusters", "KLMClusters", "MCParticles", "PIDLikelihoods"],
105 }),
106 DataStorePrinter("V0", [
107 "getTracks", "getTrackFitResults", "getV0Hypothesis",
108 "getFittedVertexX", "getFittedVertexY", "getFittedVertexZ"
109 ], {
110 "getRelationsWith": ["MCParticles"],
111 }),
112 DataStorePrinter("Kinks", [
113 "getMotherTrackIndex", "getDaughterTrackIndex", "getTrackFitResultIndexMotherStart",
114 "getTrackFitResultIndexMotherEnd", "getTrackFitResultIndexDaughter",
115 "getFittedVertexX", "getFittedVertexY", "getFittedVertexZ", "getFilterFlag"
116 ]),
117 DataStorePrinter("TrackFitResult", [
118 "getPosition", "getMomentum", "get4Momentum", "getEnergy", "getTransverseMomentum",
119 "getCovariance6", "getParticleType", "getChargeSign", "getPValue", "getD0", "getPhi0",
120 "getPhi", "getOmega", "getZ0", "getTanLambda", "getCotTheta",
121 "getTau", "getCov", "getCovariance5", "getHitPatternCDC", "getHitPatternVXD", "getNDF", "getChi2",
122 ]),
123 DataStorePrinter("EventLevelTrackingInfo", [
124 "getNCDCHitsNotAssigned", "getNCDCHitsNotAssignedPostCleaning",
125 "getNCDCSegments", "getSVDFirstSampleTime", "hasAnErrorFlag",
126 "hasUnspecifiedTrackFindingFailure", "hasVXDTF2AbortionFlag",
127 "hasSVDCKFAbortionFlag", "hasPXDCKFAbortionFlag", "hasSVDSpacePointCreatorAbortionFlag"], {
128 "hasCDCLayer": range(56)
129 }, array=False),
130 DataStorePrinter("PIDLikelihood", ["getMostLikely", "isAvailable", "areAllAvailable"], {
131 "isAvailable": pid_detectors,
132 "areAllAvailable": pid_detectors,
133 "getLogL": charged_stables,
134 "getProbability": charged_stables,
135 "getLogarithmicProbability": charged_stables,
136 }),
137 DataStorePrinter("ECLCluster", [
138 "isTrack", "isNeutral", "getStatus", "getConnectedRegionId",
139 "getClusterId", "getUniqueClusterId", "getMinTrkDistance", "getDeltaL",
140 "getAbsZernike40", "getAbsZernike51", "getZernikeMVA", "getE1oE9",
141 "getE9oE21", "getNumberOfHadronDigits", "getR", "getHypotheses",
142 "getSecondMoment", "getLAT", "getNumberOfCrystals", "getTime",
143 "getDeltaTime99", "hasFailedFitTime", "hasFailedTimeResolution", "getPhi", "getTheta",
144 "getEnergyRaw", "getEnergyHighestCrystal", "getUncertaintyEnergy",
145 "getUncertaintyTheta", "getUncertaintyPhi", "getClusterPosition",
146 "getCovarianceMatrix3x3", "getDetectorRegion",
147 "isTriggerCluster", "hasTriggerClusterMatching", "hasPulseShapeDiscrimination",
148 "getPulseShapeDiscriminationMVA", "getMaxECellId", "getMinTrkDistanceID"
149 ], {
150 "getEnergy": [16, 32],
151 "hasHypothesis": [16, 32],
152 "getRelationsWith": ["KlIds", "MCParticles"],
153 }),
154 DataStorePrinter("EventLevelClusteringInfo", [
155 "getNECLCalDigitsOutOfTimeFWD", "getNECLCalDigitsOutOfTimeBarrel",
156 "getNECLCalDigitsOutOfTimeBWD", "getNECLCalDigitsOutOfTime",
157 "getNECLShowersRejectedFWD", "getNECLShowersRejectedBarrel",
158 "getNECLShowersRejectedBWD", "getNECLShowersRejected",
159 "getNKLMDigitsMultiStripFWD", "getNKLMDigitsMultiStripBarrel",
160 "getNKLMDigitsMultiStripBWD", "getNKLMDigitsMultiStrip",
161 "getNECLShowersFWD", "getNECLShowersBarrel",
162 "getNECLShowersBWD", "getNECLShowers",
163 "getNECLLocalMaximumsFWD", "getNECLLocalMaximumsBarrel",
164 "getNECLLocalMaximumsBWD", "getNECLLocalMaximums",
165 "getNECLTriggerCellsFWD", "getNECLTriggerCellsBarrel",
166 "getNECLTriggerCellsBWD", "getNECLTriggerCells"
167 ], array=False),
168 DataStorePrinter("KLMCluster", [
169 "getTime", "getLayers", "getInnermostLayer",
170 "getClusterPosition", "getPosition", "getMomentumMag", "getEnergy",
171 "getMomentum", "getError4x4", "getError7x7",
172 "getAssociatedEclClusterFlag", "getAssociatedTrackFlag",
173 ], {
174 "getRelationsWith": ["KlIds", "MCParticles"],
175 }),
176 DataStorePrinter("KlId", ["isKLM", "isECL", "getKlId"]),
177 DataStorePrinter("TRGSummary", [
178 # getter functions to be called without arguments
179 "getTimType", "getTimQuality", "isPoissonInInjectionVeto"
180 ], {
181 # getter functions to be called with arguments
182 "getTRGSummary": range(10),
183 "getPreScale": [[int(i / 32), i % 32] for i in list(range(320))],
184 "getInputBits": range(10),
185 "getFtdlBits": range(10),
186 "getPsnmBits": range(10)
187 }, array=False),
188 DataStorePrinter("SoftwareTriggerResult", ["getResults", "getNonPrescaledResults"], array=False),
189 DataStorePrinter("MCParticle", [
190 "getPDG", "getStatus", "getMass", "getCharge", "getEnergy", "hasValidVertex",
191 "getProductionTime", "getDecayTime", "getLifetime", "getVertex",
192 "getProductionVertex", "getMomentum", "get4Vector", "getDecayVertex",
193 "getIndex", "getArrayIndex",
194 "getFirstDaughter", "getLastDaughter", "getDaughters", "getNDaughters", "getMother",
195 "getSecondaryPhysicsProcess", "getSeenInDetector",
196 "isVirtual", "isInitial", "isPrimaryParticle", "getName"
197 ], {'getDaughter': [0], 'getParticleFromGeneralizedIndexString': ['0:0']}),
198 DataStorePrinter("EventLevelTriggerTimeInfo", [
199 "isValid", "hasInjection", "isHER", "isRevo2",
200 "getTimeSinceLastInjection", "getTimeSincePrevTrigger", "getBunchNumber",
201 "getTimeSinceLastInjectionInMicroSeconds", "getTimeSincePrevTriggerInMicroSeconds",
202 "getTimeSinceInjectedBunch", "getTimeSinceInjectedBunchInMicroSeconds", "getTriggeredBunchNumberGlobal"
203 ], array=False),
204 ]
205 path.add_module(PrintObjectsModule(mdst_dataobjects, print_untested))
206