Belle II Software development
record.py
1#!/usr/bin/env python3
2
3
10
11import basf2
12
13from ROOT import Belle2 # make Belle2 namespace available
14
15import os
16import sys
17
18from tracking.validation.utilities import is_primary
19
20import tracking.harvest.harvesting as harvesting
21import tracking.harvest.refiners as refiners
22from tracking.harvest.run import HarvestingRun
23
24import trackfindingcdc.harvest.cdc_peelers as cdc_peelers
25
26import logging
27
28# @cond internal_test
29
30
31def get_logger():
32 return logging.getLogger(__name__)
33
34
35class SegmentPairFitValidationRun(HarvestingRun):
36 """Harvester to generate, postprocess and inspect MC events for track-segment-pair fit validation"""
37
38 n_events = 10000
39
40 generator_module = "simple_gun" # Rather high momentum tracks should make the tracks rather straight.
41
42
43 monte_carlo = "no"
44
45 segment_orientation = "outwards"
46
47
48 fit_method_name = "fuse-sz"
49
50 @property
51 def output_file_name(self):
52 """Get the output ROOT filename"""
53 file_name = self.fit_method_name
54 file_name += "-mc-" + self.monte_carlo
55
56 if self.root_input_file:
57 file_name += "-" + os.path.split(self.root_input_file)[1]
58 else:
59 file_name += ".root"
60
61 return file_name
62
63 def harvesting_module(self, path=None):
64 """Harvest and post-process the MC events"""
65 harvesting_module = SegmentPairFitValidationModule(self.output_file_name)
66 if path:
67 path.add_module(harvesting_module)
68 return harvesting_module
69
70 def create_argument_parser(self, **kwds):
71 """Convert command-line arguments to basf2 argument list"""
72 argument_parser = super().create_argument_parser(**kwds)
73
74 argument_parser.add_argument(
75 '-m',
76 '--monte-carlo',
77 choices=["no", "medium", "full"],
78 default=self.monte_carlo,
79 dest='monte_carlo',
80 help='Amount of monte carlo information to be used in the segment generation.',
81 )
82
83 argument_parser.add_argument(
84 "--fit",
85 choices=["zreco", "fuse-pre", "fuse-sz", "fuse-sz-re"],
86 default=self.fit_method_name,
87 dest="fit_method_name",
88 help=("Choose which fit positional information of the segment should be used. \n"
89 "* 'zreco' means the z coordinate is reconstructed and a linear sz fit is made. "
90 "No covariance between the circle and the linear sz part can be made.\n"
91 "* 'fuse-sz' means the Kalmanesk fuse of the two trajectory fits.\n"
92 "* 'fuse-sz-re' means the Kalmanesk fuse of the two trajectory fits but reestimate the drift length."
93 )
94 )
95
96 return argument_parser
97
98 def get_fit_method(self):
99 """Determine which track-segment-pair fitter to use"""
100 fit_method_name = self.fit_method_name
101
102 if fit_method_name == 'zreco':
104
105 def z_reconstruction_fit(pair):
106 return sz_fitter.update(pair)
107 return z_reconstruction_fit
108
109 elif fit_method_name.startswith('fuse-pre'):
112 fusionFit = CDCAxialStereoFusion()
113
114 def sz_segment_pair_preliminary_fit(pair):
115 fusionFit.fusePreliminary(pair)
116 return sz_segment_pair_preliminary_fit
117
118 elif fit_method_name.startswith('fuse-sz'):
121 reestimateDriftLength = fit_method_name.endswith("re")
122 fusionFit = CDCAxialStereoFusion(reestimateDriftLength)
123
124 def sz_segment_pair_fusion_fit(pair):
125 fusionFit.reconstructFuseTrajectories(pair)
126 return
127
128 trajectory3D = pair.getTrajectory3D()
129 revFromSegment = pair.getFromSegment().reversed()
130 revToSegment = pair.getToSegment().reversed()
131 revPair = CDCSegmentPair(revToSegment, revFromSegment)
132
133 CDCAxialStereoFusion.reconstructFuseTrajectories(revPair)
134 revTrajectory3D = revPair.getTrajectory3D().reversed()
135
136 # print("One origin x", trajectory3D.getLocalOrigin().x())
137 # print("One origin y", trajectory3D.getLocalOrigin().y())
138 # print("One origin z", trajectory3D.getLocalOrigin().z())
139
140 # print("Rev origin x", revTrajectory3D.getLocalOrigin().x())
141 # print("Rev origin y", revTrajectory3D.getLocalOrigin().y())
142 # print("Rev origin z", revTrajectory3D.getLocalOrigin().z())
143
144 print("One parameters", [trajectory3D.getLocalHelix().parameters()[i] for i in range(5)])
145 print("Rev parameters", [revTrajectory3D.getLocalHelix().parameters()[i] for i in range(5)])
146
147 print("One covariance")
148 for j in range(5):
149 print([trajectory3D.getLocalHelix().helixCovariance()(i, j) for i in range(5)])
150
151 print("Rev covariance")
152 for j in range(5):
153 print([revTrajectory3D.getLocalHelix().helixCovariance()(i, j) for i in range(5)])
154
155 # return revTrajectory3D
156 # return trajectory3D
157
158 return sz_segment_pair_fusion_fit
159
160 else:
161 raise ValueError(f"Unexpected fit_positions {fit_method_name}")
162
163 def create_path(self):
164 """
165 Sets up a path that plays back pregenerated events or generates events
166 based on the properties in the base class.
167 """
168 path = super().create_path()
169
170 path.add_module("TFCDC_WireHitPreparer",
171 flightTimeEstimation="outwards",
172 UseNLoops=0.5
173 )
174
175 path.add_module("TFCDC_ClusterPreparer")
176
177
178 if self.monte_carlo == "no":
179 # MC free - default
180 path.add_module("TFCDC_SegmentFinderFacetAutomaton",
181 SegmentOrientation="outwards"
182 )
183
184 path.add_module("TFCDC_SegmentFitter",
185 inputSegments="CDCSegment2DVector",
186 updateDriftLength=True,
187 useAlphaInDriftLength=True,
188 )
189
190 elif self.monte_carlo == "medium":
191 # Medium MC - proper generation logic, but true facets and facet relations
192 path.add_module("TFCDC_SegmentFinderFacetAutomaton",
193 FacetFilter="truth",
194 FacetRelationFilter="truth",
195 SegmentOrientation="outwards"
196 )
197
198 path.add_module("TFCDC_SegmentFitter",
199 inputSegments="CDCSegment2DVector",
200 updateDriftLength=True,
201 useAlphaInDriftLength=True,
202 )
203
204 elif self.monte_carlo == "full":
205 # Only true monte carlo segments
206 # make the positions realistic but keep the true drift length
207 path.add_module("TFCDC_SegmentCreatorMCTruth",
208 reconstructedDriftLength=False,
209 reconstructedPositions=True,
210 # segments="MCSegments"
211 )
212
213 path.add_module("TFCDC_SegmentFitter",
214 inputSegments="CDCSegment2DVector",
215 updateDriftLength=False,
216 # useAlphaInDriftLength=True,
217 )
218
219 else:
220 raise ValueError("Invalid degree of Monte Carlo information")
221
222 path.add_module("TFCDC_SegmentOrienter",
223 SegmentOrientation="outwards",
224 # SegmentOrientation="none",
225 inputSegments="CDCSegment2DVector",
226 segments="CDCSegment2DVectorOriented"
227 )
228
229 path.add_module("TFCDC_TrackFinderSegmentPairAutomaton",
230 inputSegments="CDCSegment2DVectorOriented",
231 WriteSegmentPairs=True,
232 SegmentPairFilter="truth",
233 SegmentPairFilterParameters={"allowReverse": True},
234 SegmentPairRelationFilter="none"
235 )
236
237 path.add_module(AxialStereoPairFitterModule(fit_method=self.get_fit_method()))
238 return path
239
240
241class SegmentPairFitValidationModule(harvesting.HarvestingModule):
242
243 """Module to collect information about the generated segment pairs and
244 compose validation plots on terminate."""
245
246 def __init__(self, output_file_name):
247 """Constructor"""
248 super().__init__(
249 output_file_name=output_file_name,
250 foreach="CDCSegmentPairVector"
251 )
252
253 ## by default, there is no method to find matching MC track segments
254 self.mc_segment_lookup = None
255
256 def initialize(self):
257 """Receive signal at the start of event processing"""
258 self.mc_segment_lookup = Belle2.TrackFindingCDC.CDCMCSegment2DLookUp.getInstance()
259 super().initialize()
260
261 def prepare(self):
262 """Initialize the MC-hit lookup method"""
263 Belle2.TrackFindingCDC.CDCMCHitLookUp.getInstance().fill()
264
265 def pick(self, segment_pair):
266 """Select segment pairs with 4 or more hits per segment and a matching primary MC particle"""
267 mc_segment_lookup = self.mc_segment_lookup
268 from_segment = segment_pair.getFromSegment()
269 to_segment = segment_pair.getToSegment()
270 mc_particle = mc_segment_lookup.getMCParticle(from_segment)
271 return (mc_particle and
272 is_primary(mc_particle) and
273 from_segment.size() > 3 and
274 to_segment.size() > 3)
275
276 def peel(self, segment_pair):
277 """Aggregate the track and MC information for track-segment analysis"""
278 mc_segment_lookup = self.mc_segment_lookup
279
280 # from_segment = segment_pair.getFromSegment()
281 to_segment = segment_pair.getToSegment()
282
283 # mc_particle = mc_segment_lookup.getMCParticle(from_segment)
284
285 # Take the fit best at the middle of the segment pair
286 fit3d_truth = mc_segment_lookup.getTrajectory3D(to_segment)
287
288 fb_info = 1 if segment_pair.getAutomatonCell().getCellWeight() > 0 else -1
289 truth_crops = dict(
290 curvature_truth=fb_info * fit3d_truth.getCurvatureXY(),
291 tan_lambda_truth=fb_info * fit3d_truth.getTanLambda(),
292 )
293
294 segment_pair_crops = cdc_peelers.peel_segment_pair(segment_pair)
295 segment_pair_crops.update(truth_crops)
296 return segment_pair_crops
297
298 # Refiners to be executed at the end of the harvesting / termination of the module
299 ## Save histograms in a sub folder
300 save_histograms = refiners.save_histograms(outlier_z_score=5.0, allow_discrete=True)
301 ## Save a tree of all collected variables in a sub folder
302 save_tree = refiners.save_tree()
303
304 ## Save curvature-pull auxiliary information in a sub folder
305 save_curvature_pull_aux = refiners.save_pull_analysis(
306 part_name="curvature",
307 folder_name="aux",
308 unit="1/cm",
309 absolute=False,
310 aux_names=["superlayer_id_pair", "tan_lambda_truth"],
311 which_plots=["aux"],
312 outlier_z_score=3.0)
313
314 ## Save curvature-pull information in a sub folder
315 save_curvature_pull = refiners.save_pull_analysis(
316 part_name="curvature",
317 unit="1/cm",
318 absolute=False,
319 groupby=[None, "superlayer_id_pair"],
320 outlier_z_score=3.0)
321
322 ## Save absolute curvature-pull information in a sub folder
323 save_absolute_curvature_pull = refiners.save_pull_analysis(
324 part_name="curvature",
325 unit="1/cm",
326 absolute=True,
327 groupby=[None, "superlayer_id_pair"],
328 outlier_z_score=3.0)
329
330 ## Save tan(lambda)-fit pull information in a sub folder
331 save_tan_lambda_pull = refiners.save_pull_analysis(
332 part_name="tan_lambda",
333 quantity_name="tan #lambda",
334 absolute=False,
335 groupby=[None, "superlayer_id_pair"],
336 outlier_z_score=3.0)
337
338 ## Save fit-quality histograms in a sub folder
339 save_fit_quality_histograms = refiners.save_histograms(
340 outlier_z_score=5.0,
341 select={"ndf": "ndf",
342 "chi2": "#chi2",
343 "p_value": "p-value",
344 },
345 groupby=[None, "superlayer_id_pair"],
346 description="Distribution of {part_name} in the segment fits",
347 )
348
349 ## Save tan(lambda)-fit-quality profile histograms in a sub folder
350 save_fit_quality_by_tan_lambda_profiles = refiners.save_profiles(
351 select={
352 "p_value": "fit p-value",
353 "tan_lambda_truth": "true tan #lambda",
354 },
355 groupby=[None, "superlayer_id_pair"],
356 x="true tan #lambda",
357 y="fit p-value",
358 check=("The {y_part_name} should be essentially the same over"
359 "the whole range of the tan lambda spectrum"),
360 description=("Investigating the reconstruction quality for different "
361 "tan lambda regions of the CDC. Most notably is the superlayer dependency."
362 "For stereo superlayers the curve is not flat but has distinct slope."),
363 fit='line',
364 )
365
366
367class AxialStereoPairFitterModule(basf2.Module):
368 """Fit axial and stereo track segment pairs"""
369
370 @staticmethod
371 def default_fit_method(segmentPair):
372 """Default method to fit the generated segment pairs."""
373
374 CDCAxialStereoFusion = Belle2.TrackFindingCDC.CDCAxialStereoFusion
375 CDCAxialStereoFusion.reconstructFuseTrajectories(segmentPair,
376 True)
377
378 def __init__(self, fit_method=None):
379 """Constructor"""
380
381 ## fit_method : function
382 ## A function called on each stored segment pair as its only argument to update its fit inplace.
383 ## See default_fit_method for an example. Defaults to None meaning the default_fit_method is used
384 ## Method used to fit the individual segment pairs
385 self.fit_method = fit_method
386 if not fit_method:
387 self.fit_method = self.default_fit_method
388
389 super().__init__()
390
391 def event(self):
392 """Called by basf2 for each event"""
393 self.fitStoredPairs()
394
395 def fitStoredPairs(self):
396 """Fits all pairs in the StoreArray with designated fit method."""
397
398 fit_method = self.fit_method
399
400 stored_segment_pair_relations = Belle2.PyStoreObj("CDCSegmentPairVector")
401 wrapped_segment_pair_relations = stored_segment_pair_relations.obj()
402 segment_pair_relations = wrapped_segment_pair_relations.get()
403
404 for segment_pair_relation in segment_pair_relations:
405 fit_method(segment_pair_relation)
406
407
408def main():
409 run = SegmentPairFitValidationRun()
410 run.configure_and_execute_from_commandline()
411
412
413if __name__ == "__main__":
414 logging.basicConfig(stream=sys.stdout, level=logging.INFO, format='%(levelname)s:%(message)s')
415 main()
416
417# @endcond
Utility class implementing the Kalmanesk combination of to two dimensional trajectories to one three ...
static const CDCSZFitter & getFitter()
Getter for a standard sz line fitter instance.
Class representing a pair of one reconstructed axial segment and one stereo segment in adjacent super...