refactor: python reconstruction trajectory writing (#1718)

rearrange writing mechanism for trajectories and reuse it for ambires and others in the future

Examples/Python/python/acts/examples/reconstruction.py

@@ -129,6 +129,24 @@
     defaults=[(None, None)] * 7 + [None] * 2,
 )
 
+CKFPerformanceConfig = namedtuple(
+    "CKFPerformanceConfig",
+    ["truthMatchProbMin", "nMeasurementsMin", "ptMin"],
+    defaults=[None] * 3,
+)
+
+AmbiguityResolutionConfig = namedtuple(
+    "AmbiguityResolutionConfig",
+    ["maximumSharedHits"],
+    defaults=[None] * 1,
+)
+
+
+class VertexFinder(Enum):
+    Truth = (1,)
+    AMVF = (2,)
+    Iterative = (3,)
+
 
 @acts.examples.NamedTypeArgs(
     seedingAlgorithm=SeedingAlgorithm,
@@ -676,26 +694,19 @@ def addTruthTrackingGsf(
     return s
 
 
-CKFPerformanceConfig = namedtuple(
-    "CKFPerformanceConfig",
-    ["truthMatchProbMin", "nMeasurementsMin", "ptMin"],
-    defaults=[None] * 3,
-)
-
-
 @acts.examples.NamedTypeArgs(
-    CKFPerformanceConfigArg=CKFPerformanceConfig,
+    ckfPerformanceConfig=CKFPerformanceConfig,
     trackSelectorRanges=TrackSelectorRanges,
 )
 def addCKFTracks(
     s: acts.examples.Sequencer,
     trackingGeometry: acts.TrackingGeometry,
     field: acts.MagneticFieldProvider,
-    CKFPerformanceConfigArg: CKFPerformanceConfig = CKFPerformanceConfig(),
+    trackSelectorRanges: Optional[TrackSelectorRanges] = None,
+    selectedParticles: str = "truth_seeds_selected",
+    ckfPerformanceConfig: CKFPerformanceConfig = CKFPerformanceConfig(),
     outputDirCsv: Optional[Union[Path, str]] = None,
     outputDirRoot: Optional[Union[Path, str]] = None,
-    selectedParticles: str = "truth_seeds_selected",
-    trackSelectorRanges: Optional[TrackSelectorRanges] = None,
     writeTrajectories: bool = True,
     logLevel: Optional[acts.logging.Level] = None,
 ) -> None:
@@ -707,7 +718,7 @@ def addCKFTracks(
         the sequencer module to which we add the Seeding steps (returned from addSeeding)
     trackingGeometry : tracking geometry
     field : magnetic field
-    CKFPerformanceConfigArg : CKFPerformanceConfig(truthMatchProbMin, nMeasurementsMin, ptMin)
+    ckfPerformanceConfigArg : CKFPerformanceConfig(truthMatchProbMin, nMeasurementsMin, ptMin)
         CKFPerformanceWriter configuration.
         Defaults specified in Examples/Io/Performance/ActsExamples/Io/Performance/CKFPerformanceWriter.hpp
     outputDirCsv : Path|str, path, None
@@ -724,7 +735,6 @@ def addCKFTracks(
     """
 
     customLogLevel = acts.examples.defaultLogging(s, logLevel)
-    logger = acts.logging.getLogger("addCKFTracks")
 
     # Setup the track finding algorithm with CKF
     # It takes all the source links created from truth hit smearing, seeds from
@@ -757,16 +767,54 @@ def addCKFTracks(
 
         s.addWhiteboardAlias("trajectories", trackSelector.config.outputTrajectories)
 
+    addTrajectoryWriters(
+        s,
+        name="ckf",
+        trajectories=trackFinder.config.outputTrajectories,
+        ckfPerformanceConfig=ckfPerformanceConfig,
+        outputDirCsv=outputDirCsv,
+        outputDirRoot=outputDirRoot,
+        writeStates=writeTrajectories,
+        writeSummary=writeTrajectories,
+        writeCKFperformance=True,
+        writeFinderPerformance=False,
+        writeFitterPerformance=False,
+        logLevel=logLevel,
+    )
+
+    return s
+
+
+@acts.examples.NamedTypeArgs(
+    ckfPerformanceConfig=CKFPerformanceConfig,
+    trackSelectorRanges=TrackSelectorRanges,
+)
+def addTrajectoryWriters(
+    s: acts.examples.Sequencer,
+    name: str,
+    trajectories: str = "trajectories",
+    ckfPerformanceConfig: CKFPerformanceConfig = CKFPerformanceConfig(),
+    outputDirCsv: Optional[Union[Path, str]] = None,
+    outputDirRoot: Optional[Union[Path, str]] = None,
+    writeStates: bool = True,
+    writeSummary: bool = True,
+    writeCKFperformance: bool = True,
+    writeFinderPerformance: bool = True,
+    writeFitterPerformance: bool = True,
+    logLevel: Optional[acts.logging.Level] = None,
+):
+    customLogLevel = acts.examples.defaultLogging(s, logLevel)
+
     if outputDirRoot is not None:
         outputDirRoot = Path(outputDirRoot)
         if not outputDirRoot.exists():
             outputDirRoot.mkdir()
 
-        if writeTrajectories:
+        if writeStates:
             # write track states from CKF
             trackStatesWriter = acts.examples.RootTrajectoryStatesWriter(
                 level=customLogLevel(),
-                inputTrajectories=trackFinder.config.outputTrajectories,
+                inputTrajectories=trajectories,
                 # @note The full particles collection is used here to avoid lots of warnings
                 # since the unselected CKF track might have a majority particle not in the
                 # filtered particle collection. This could be avoided when a seperate track
@@ -775,56 +823,83 @@ def addCKFTracks(
                 inputSimHits="simhits",
                 inputMeasurementParticlesMap="measurement_particles_map",
                 inputMeasurementSimHitsMap="measurement_simhits_map",
-                filePath=str(outputDirRoot / "trackstates_ckf.root"),
+                filePath=str(outputDirRoot / f"trackstates_{name}.root"),
                 treeName="trackstates",
             )
             s.addWriter(trackStatesWriter)
 
+        if writeSummary:
             # write track summary from CKF
             trackSummaryWriter = acts.examples.RootTrajectorySummaryWriter(
                 level=customLogLevel(),
-                inputTrajectories=trackFinder.config.outputTrajectories,
+                inputTrajectories=trajectories,
                 # @note The full particles collection is used here to avoid lots of warnings
                 # since the unselected CKF track might have a majority particle not in the
                 # filtered particle collection. This could be avoided when a seperate track
                 # selection algorithm is used.
                 inputParticles="particles_selected",
                 inputMeasurementParticlesMap="measurement_particles_map",
-                filePath=str(outputDirRoot / "tracksummary_ckf.root"),
+                filePath=str(outputDirRoot / f"tracksummary_{name}.root"),
                 treeName="tracksummary",
             )
             s.addWriter(trackSummaryWriter)
 
-        # Write CKF performance data
-        ckfPerfWriter = acts.examples.CKFPerformanceWriter(
-            level=customLogLevel(),
-            inputParticles=selectedParticles,
-            inputTrajectories=trackFinder.config.outputTrajectories,
-            inputMeasurementParticlesMap="measurement_particles_map",
-            **acts.examples.defaultKWArgs(
-                # The bottom seed could be the first, second or third hits on the truth track
-                nMeasurementsMin=CKFPerformanceConfigArg.nMeasurementsMin,
-                ptMin=CKFPerformanceConfigArg.ptMin,
-                truthMatchProbMin=CKFPerformanceConfigArg.truthMatchProbMin,
-            ),
-            filePath=str(outputDirRoot / "performance_ckf.root"),
-        )
-        s.addWriter(ckfPerfWriter)
+        if writeCKFperformance:
+            # Write CKF performance data
+            ckfPerfWriter = acts.examples.CKFPerformanceWriter(
+                level=customLogLevel(),
+                inputParticles="truth_seeds_selected",
+                inputTrajectories=trajectories,
+                inputMeasurementParticlesMap="measurement_particles_map",
+                **acts.examples.defaultKWArgs(
+                    # The bottom seed could be the first, second or third hits on the truth track
+                    nMeasurementsMin=ckfPerformanceConfig.nMeasurementsMin,
+                    ptMin=ckfPerformanceConfig.ptMin,
+                    truthMatchProbMin=ckfPerformanceConfig.truthMatchProbMin,
+                ),
+                filePath=str(outputDirRoot / f"performance_{name}.root"),
+            )
+            s.addWriter(ckfPerfWriter)
+
+        if writeFinderPerformance:
+            s.addWriter(
+                acts.examples.TrackFinderPerformanceWriter(
+                    level=acts.logging.INFO,
+                    inputProtoTracks="prototracks",
+                    inputParticles="truth_seeds_selected",
+                    inputMeasurementParticlesMap="measurement_particles_map",
+                    filePath=str(
+                        outputDirRoot / f"performance_track_finder_{name}.root"
+                    ),
+                )
+            )
+
+        if writeFitterPerformance:
+            s.addWriter(
+                acts.examples.TrackFitterPerformanceWriter(
+                    level=acts.logging.INFO,
+                    inputTrajectories="trajectories",
+                    inputParticles="truth_seeds_selected",
+                    inputMeasurementParticlesMap="measurement_particles_map",
+                    filePath=str(
+                        outputDirRoot / f"performance_track_fitter_{name}.root"
+                    ),
+                )
+            )
 
     if outputDirCsv is not None:
         outputDirCsv = Path(outputDirCsv)
         if not outputDirCsv.exists():
             outputDirCsv.mkdir()
-        logger.info("Writing CSV files")
-        csvMTJWriter = acts.examples.CsvMultiTrajectoryWriter(
-            level=customLogLevel(),
-            inputTrajectories=trackFinder.config.outputTrajectories,
-            inputMeasurementParticlesMap="measurement_particles_map",
-            outputDir=str(outputDirCsv),
-        )
-        s.addWriter(csvMTJWriter)
 
-    return s
+        if trackSummaryWriter:
+            csvMTJWriter = acts.examples.CsvMultiTrajectoryWriter(
+                level=customLogLevel(),
+                inputTrajectories=trajectories,
+                inputMeasurementParticlesMap="measurement_particles_map",
+                outputDir=str(outputDirCsv),
+            )
+            s.addWriter(csvMTJWriter)
 
 
 @acts.examples.NamedTypeArgs(
@@ -961,22 +1036,18 @@ def addExaTrkX(
     return s
 
 
-AmbiguityResolutionConfig = namedtuple(
-    "AmbiguityResolutionConfig",
-    ["maximumSharedHits"],
-    defaults=[None] * 1,
-)
-
-
 @acts.examples.NamedTypeArgs(
     config=AmbiguityResolutionConfig,
-    ckfPerformanceConfigArg=CKFPerformanceConfig,
+    ckfPerformanceConfig=CKFPerformanceConfig,
 )
 def addAmbiguityResolution(
     s,
     config: AmbiguityResolutionConfig = AmbiguityResolutionConfig(),
-    ckfPerformanceConfigArg: CKFPerformanceConfig = CKFPerformanceConfig(),
+    selectedParticles: str = "truth_seeds_selected",
+    ckfPerformanceConfig: CKFPerformanceConfig = CKFPerformanceConfig(),
+    outputDirCsv: Optional[Union[Path, str]] = None,
     outputDirRoot: Optional[Union[Path, str]] = None,
+    writeTrajectories: bool = True,
     logLevel: Optional[acts.logging.Level] = None,
 ) -> None:
 
@@ -997,34 +1068,24 @@ def addAmbiguityResolution(
 
     s.addWhiteboardAlias("trajectories", alg.config.outputTrajectories)
 
-    if outputDirRoot is not None:
-        outputDirRoot = Path(outputDirRoot)
-        if not outputDirRoot.exists():
-            outputDirRoot.mkdir()
-
-        perfWriter = acts.examples.CKFPerformanceWriter(
-            level=customLogLevel(),
-            inputParticles="truth_seeds_selected",
-            inputTrajectories=alg.config.inputTrajectories,
-            inputMeasurementParticlesMap="measurement_particles_map",
-            **acts.examples.defaultKWArgs(
-                nMeasurementsMin=ckfPerformanceConfigArg.nMeasurementsMin,
-                ptMin=ckfPerformanceConfigArg.ptMin,
-                truthMatchProbMin=ckfPerformanceConfigArg.truthMatchProbMin,
-            ),
-            filePath=str(outputDirRoot / "performance_ambi.root"),
-        )
-        s.addWriter(perfWriter)
+    addTrajectoryWriters(
+        s,
+        name="ambi",
+        trajectories=alg.config.outputTrajectories,
+        ckfPerformanceConfig=ckfPerformanceConfig,
+        outputDirCsv=outputDirCsv,
+        outputDirRoot=outputDirRoot,
+        writeStates=writeTrajectories,
+        writeSummary=writeTrajectories,
+        writeCKFperformance=True,
+        writeFinderPerformance=False,
+        writeFitterPerformance=False,
+        logLevel=logLevel,
+    )
 
     return s
 
 
-class VertexFinder(Enum):
-    Truth = (1,)
-    AMVF = (2,)
-    Iterative = (3,)
-
-
 @acts.examples.NamedTypeArgs(
     trackSelectorRanges=TrackSelectorRanges,
 )