13from ROOT
import Belle2
24import trackfindingcdc.harvest.cdc_peelers
as cdc_peelers
32 return logging.getLogger(__name__)
35class SegmentPairFitValidationRun(HarvestingRun):
36 """Harvester to generate, postprocess and inspect MC events for track-segment-pair fit validation"""
40 generator_module =
"simple_gun"
45 segment_orientation =
"outwards"
48 fit_method_name =
"fuse-sz"
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
56 if self.root_input_file:
57 file_name +=
"-" + os.path.split(self.root_input_file)[1]
63 def harvesting_module(self, path=None):
64 """Harvest and post-process the MC events"""
65 harvesting_module = SegmentPairFitValidationModule(self.output_file_name)
67 path.add_module(harvesting_module)
68 return harvesting_module
70 def create_argument_parser(self, **kwds):
71 """Convert command-line arguments to basf2 argument list"""
72 argument_parser = super().create_argument_parser(**kwds)
74 argument_parser.add_argument(
77 choices=[
"no",
"medium",
"full"],
78 default=self.monte_carlo,
80 help=
'Amount of monte carlo information to be used in the segment generation.',
83 argument_parser.add_argument(
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."
96 return argument_parser
98 def get_fit_method(self):
99 """Determine which track-segment-pair fitter to use"""
100 fit_method_name = self.fit_method_name
102 if fit_method_name ==
'zreco':
105 def z_reconstruction_fit(pair):
106 return sz_fitter.update(pair)
107 return z_reconstruction_fit
109 elif fit_method_name.startswith(
'fuse-pre'):
112 fusionFit = CDCAxialStereoFusion()
114 def sz_segment_pair_preliminary_fit(pair):
115 fusionFit.fusePreliminary(pair)
116 return sz_segment_pair_preliminary_fit
118 elif fit_method_name.startswith(
'fuse-sz'):
121 reestimateDriftLength = fit_method_name.endswith(
"re")
122 fusionFit = CDCAxialStereoFusion(reestimateDriftLength)
124 def sz_segment_pair_fusion_fit(pair):
125 fusionFit.reconstructFuseTrajectories(pair)
128 trajectory3D = pair.getTrajectory3D()
129 revFromSegment = pair.getFromSegment().reversed()
130 revToSegment = pair.getToSegment().reversed()
131 revPair = CDCSegmentPair(revToSegment, revFromSegment)
133 CDCAxialStereoFusion.reconstructFuseTrajectories(revPair)
134 revTrajectory3D = revPair.getTrajectory3D().reversed()
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)])
147 print(
"One covariance")
149 print([trajectory3D.getLocalHelix().helixCovariance()(i, j)
for i
in range(5)])
151 print(
"Rev covariance")
153 print([revTrajectory3D.getLocalHelix().helixCovariance()(i, j)
for i
in range(5)])
158 return sz_segment_pair_fusion_fit
161 raise ValueError(f
"Unexpected fit_positions {fit_method_name}")
163 def create_path(self):
165 Sets up a path that plays back pregenerated events or generates events
166 based on the properties in the base class.
168 path = super().create_path()
170 path.add_module(
"TFCDC_WireHitPreparer",
171 flightTimeEstimation=
"outwards",
175 path.add_module(
"TFCDC_ClusterPreparer")
178 if self.monte_carlo ==
"no":
180 path.add_module(
"TFCDC_SegmentFinderFacetAutomaton",
181 SegmentOrientation=
"outwards"
184 path.add_module(
"TFCDC_SegmentFitter",
185 inputSegments=
"CDCSegment2DVector",
186 updateDriftLength=
True,
187 useAlphaInDriftLength=
True,
190 elif self.monte_carlo ==
"medium":
192 path.add_module(
"TFCDC_SegmentFinderFacetAutomaton",
194 FacetRelationFilter=
"truth",
195 SegmentOrientation=
"outwards"
198 path.add_module(
"TFCDC_SegmentFitter",
199 inputSegments=
"CDCSegment2DVector",
200 updateDriftLength=
True,
201 useAlphaInDriftLength=
True,
204 elif self.monte_carlo ==
"full":
207 path.add_module(
"TFCDC_SegmentCreatorMCTruth",
208 reconstructedDriftLength=
False,
209 reconstructedPositions=
True,
213 path.add_module(
"TFCDC_SegmentFitter",
214 inputSegments=
"CDCSegment2DVector",
215 updateDriftLength=
False,
220 raise ValueError(
"Invalid degree of Monte Carlo information")
222 path.add_module(
"TFCDC_SegmentOrienter",
223 SegmentOrientation=
"outwards",
225 inputSegments=
"CDCSegment2DVector",
226 segments=
"CDCSegment2DVectorOriented"
229 path.add_module(
"TFCDC_TrackFinderSegmentPairAutomaton",
230 inputSegments=
"CDCSegment2DVectorOriented",
231 WriteSegmentPairs=
True,
232 SegmentPairFilter=
"truth",
233 SegmentPairFilterParameters={
"allowReverse":
True},
234 SegmentPairRelationFilter=
"none"
237 path.add_module(AxialStereoPairFitterModule(fit_method=self.get_fit_method()))
241class SegmentPairFitValidationModule(harvesting.HarvestingModule):
243 """Module to collect information about the generated segment pairs and
244 compose validation plots on terminate."""
246 def __init__(self, output_file_name):
249 output_file_name=output_file_name,
250 foreach="CDCSegmentPairVector"
253 ## by default, there is no method to find matching MC track segments
254 self.mc_segment_lookup = None
256 def initialize(self):
257 """Receive signal at the start of event processing"""
258 self.mc_segment_lookup = Belle2.TrackFindingCDC.CDCMCSegment2DLookUp.getInstance()
262 """Initialize the MC-hit lookup method"""
263 Belle2.TrackFindingCDC.CDCMCHitLookUp.getInstance().fill()
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)
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
280 # from_segment = segment_pair.getFromSegment()
281 to_segment = segment_pair.getToSegment()
283 # mc_particle = mc_segment_lookup.getMCParticle(from_segment)
285 # Take the fit best at the middle of the segment pair
286 fit3d_truth = mc_segment_lookup.getTrajectory3D(to_segment)
288 fb_info = 1 if segment_pair.getAutomatonCell().getCellWeight() > 0 else -1
290 curvature_truth=fb_info * fit3d_truth.getCurvatureXY(),
291 tan_lambda_truth=fb_info * fit3d_truth.getTanLambda(),
294 segment_pair_crops = cdc_peelers.peel_segment_pair(segment_pair)
295 segment_pair_crops.update(truth_crops)
296 return segment_pair_crops
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()
304 ## Save curvature-pull auxiliary information in a sub folder
305 save_curvature_pull_aux = refiners.save_pull_analysis(
306 part_name="curvature",
310 aux_names=["superlayer_id_pair", "tan_lambda_truth"],
314 ## Save curvature-pull information in a sub folder
315 save_curvature_pull = refiners.save_pull_analysis(
316 part_name="curvature",
319 groupby=[None, "superlayer_id_pair"],
322 ## Save absolute curvature-pull information in a sub folder
323 save_absolute_curvature_pull = refiners.save_pull_analysis(
324 part_name="curvature",
327 groupby=[None, "superlayer_id_pair"],
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",
335 groupby=[None, "superlayer_id_pair"],
338 ## Save fit-quality histograms in a sub folder
339 save_fit_quality_histograms = refiners.save_histograms(
341 select={"ndf": "ndf",
343 "p_value": "p-value",
345 groupby=[None, "superlayer_id_pair"],
346 description="Distribution of {part_name} in the segment fits",
349 ## Save tan(lambda)-fit-quality profile histograms in a sub folder
350 save_fit_quality_by_tan_lambda_profiles = refiners.save_profiles(
352 "p_value": "fit p-value",
353 "tan_lambda_truth": "true tan #lambda",
355 groupby=[None, "superlayer_id_pair"],
356 x="true tan #lambda",
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."),
367class AxialStereoPairFitterModule(basf2.Module):
368 """Fit axial and stereo track segment pairs"""
371 def default_fit_method(segmentPair):
372 """Default method to fit the generated segment pairs."""
374 CDCAxialStereoFusion = Belle2.TrackFindingCDC.CDCAxialStereoFusion
375 CDCAxialStereoFusion.reconstructFuseTrajectories(segmentPair,
378 def __init__(self, fit_method=None):
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
387 self.fit_method = self.default_fit_method
392 """Called by basf2 for each event"""
393 self.fitStoredPairs()
395 def fitStoredPairs(self):
396 """Fits all pairs in the StoreArray with designated fit method."""
398 fit_method = self.fit_method
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()
404 for segment_pair_relation in segment_pair_relations:
405 fit_method(segment_pair_relation)
409 run = SegmentPairFitValidationRun()
410 run.configure_and_execute_from_commandline()
413if __name__ == "__main__":
414 logging.basicConfig(stream=sys.stdout, level=logging.INFO, format='%(levelname)s:%(message)s')
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...