feat: examples restore vertexing track selector (#1644)

#1579 remove the track selector from `addVertexing` which is still required. see #1636 #1588 should clean this up a little
acts-project · Nov 1, 2022 · e1b8ebb · e1b8ebb
1 parent 65ea889
commit e1b8ebb
Showing 1 changed file with 21 additions and 5 deletions.
diff --git a/Examples/Python/python/acts/examples/reconstruction.py b/Examples/Python/python/acts/examples/reconstruction.py
@@ -681,7 +681,7 @@ def addCKFTracks(
     outputDirCsv: Optional[Union[Path, str]] = None,
     outputDirRoot: Optional[Union[Path, str]] = None,
     selectedParticles: str = "truth_seeds_selected",
-    trackSelectorRanges: Optional[TrackSelectorRanges] = TrackSelectorRanges(),
+    trackSelectorRanges: Optional[TrackSelectorRanges] = None,
     writeTrajectories: bool = True,
     logLevel: Optional[acts.logging.Level] = None,
 ) -> None:
@@ -728,11 +728,9 @@ def addCKFTracks(
         inputInitialTrackParameters="estimatedparameters",
         outputTrajectories="trajectories",
         outputTrackParameters=outputTrackParameters
-        if trackSelectorRanges is None
-        else "fittedTrackParametersTmp",
+        + ("" if trackSelectorRanges is None else "Tmp"),
         outputTrackParametersTips=outputTrackParametersTips
-        if trackSelectorRanges is None
-        else "fittedTrackParametersTipsTmp",
+        + ("" if trackSelectorRanges is None else "Tmp"),
         findTracks=acts.examples.TrackFindingAlgorithm.makeTrackFinderFunction(
             trackingGeometry, field
         ),
@@ -1016,6 +1014,9 @@ class VertexFinder(Enum):
     Iterative = (3,)
 
 
+@acts.examples.NamedTypeArgs(
+    trackSelectorRanges=TrackSelectorRanges,
+)
 def addVertexFitting(
     s,
     field,
@@ -1025,6 +1026,7 @@ def addVertexFitting(
     trackParameters: str = "filteredTrackParameters",
     trackParametersTips: Optional[str] = "filteredTrackParametersTips",
     vertexFinder: VertexFinder = VertexFinder.Truth,
+    trackSelectorRanges: Optional[TrackSelectorRanges] = None,
     logLevel: Optional[acts.logging.Level] = None,
 ) -> None:
     """This function steers the vertex fitting
@@ -1053,6 +1055,20 @@ def addVertexFitting(
 
     customLogLevel = acts.examples.defaultLogging(s, logLevel)
 
+    if trackSelectorRanges is not None:
+        addTrackSelection(
+            s,
+            trackSelectorRanges,
+            inputTrackParameters=trackParameters,
+            inputTrackParametersTips=trackParametersTips,
+            outputTrackParameters=trackParameters + "Tmp",
+            outputTrackParametersTips=trackParametersTips + "Tmp",
+            logLevel=customLogLevel(),
+        )
+
+        trackParameters = trackParameters + "Tmp"
+        trackParametersTips = trackParametersTips + "Tmp"
+
     inputParticles = "particles_input"
     outputVertices = "fittedVertices"
     selectedParticles = "particles_selected"