refactor: use TrackSelector in addVertexFitting (#1545)

follow-up of #1538

@timadye noticed that we have to use `s.addAlgorithm(acts.examples.TrackSelector(...` in the full chain examples in order to have vertexing. this PR integrates the track selection into the `addVertexFitting` helper

CI/physmon/physmon.py

@@ -40,6 +40,7 @@
     CKFPerformanceConfig,
     addVertexFitting,
     VertexFinder,
+    TrackSelectorRanges,
 )
 
 
@@ -177,22 +178,15 @@
             outputDirCsv=None,
         )
 
-        s.addAlgorithm(
-            acts.examples.TrackSelector(
-                level=acts.logging.INFO,
-                inputTrackParameters="fittedTrackParameters",
-                outputTrackParameters="trackparameters",
-                outputTrackIndices="outputTrackIndices",
-                removeNeutral=True,
-                absEtaMax=2.5,
-                loc0Max=4.0 * u.mm,  # rho max
-                ptMin=500 * u.MeV,
-            )
-        )
-
-        s = addVertexFitting(
+        addVertexFitting(
             s,
             field,
+            TrackSelectorRanges(
+                removeNeutral=True,
+                absEta=(None, 2.5),
+                loc0=(None, 4.0 * u.mm),
+                pt=(500 * u.MeV, None),
+            ),
             vertexFinder=VertexFinder.Iterative,
             trajectories="trajectories",
             outputDirRoot=tp,

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

@@ -111,6 +111,22 @@
     defaults=[(None, None)],
 )
 
+TrackSelectorRanges = namedtuple(
+    "TrackSelectorRanges",
+    [
+        "loc0",
+        "loc1",
+        "time",
+        "eta",
+        "absEta",
+        "pt",
+        "phi",
+        "removeNeutral",
+        "removeCharged",
+    ],
+    defaults=[(None, None)] * 7 + [None] * 2,
+)
+
 
 @acts.examples.NamedTypeArgs(
     seedingAlgorithm=SeedingAlgorithm,
@@ -884,14 +900,17 @@ class VertexFinder(Enum):
     Iterative = (3,)
 
 
+@acts.examples.NamedTypeArgs(trackSelectorRanges=TrackSelectorRanges)
 def addVertexFitting(
     s,
     field,
     outputDirRoot: Optional[Union[Path, str]] = None,
     associatedParticles: str = "particles_input",
     trajectories: Optional[str] = None,
-    trackParameters: str = "trackparameters",
+    trackParameters: str = "fittedTrackParameters",
+    trackIndices: str = "outputTrackIndices",
     vertexFinder: VertexFinder = VertexFinder.Truth,
+    trackSelectorRanges: TrackSelectorRanges = TrackSelectorRanges(),
     logLevel: Optional[acts.logging.Level] = None,
 ) -> None:
     """This function steers the vertex fitting
@@ -907,6 +926,9 @@ def addVertexFitting(
         RootVertexPerformanceWriter.inputAssociatedTruthParticles
     vertexFinder : VertexFinder, Truth
         vertexFinder algorithm: one of Truth, AMVF, Iterative
+    trackSelectorRanges : TrackSelectorRanges(loc0, loc1, time, eta, absEta, pt, phi, removeNeutral, removeCharged)
+        TrackSelector configuration. Each range is specified as a tuple of (min,max).
+        Defaults of no cuts specified in Examples/Algorithms/TruthTracking/ActsExamples/TruthTracking/TrackSelector.hpp
     logLevel : acts.logging.Level, None
         logging level to override setting given in `s`
     """
@@ -920,6 +942,37 @@ def addVertexFitting(
 
     customLogLevel = acts.examples.defaultLogging(s, logLevel)
 
+    if trackSelectorRanges is not None:
+        selectedTrackParameters = "trackparameters"
+
+        trackSelector = acts.examples.TrackSelector(
+            level=customLogLevel(),
+            inputTrackParameters=trackParameters,
+            outputTrackParameters=selectedTrackParameters,
+            outputTrackIndices=trackIndices,
+            **acts.examples.defaultKWArgs(
+                loc0Min=trackSelectorRanges.loc0[0],
+                loc0Max=trackSelectorRanges.loc0[1],
+                loc1Min=trackSelectorRanges.loc1[0],
+                loc1Max=trackSelectorRanges.loc1[1],
+                timeMin=trackSelectorRanges.time[0],
+                timeMax=trackSelectorRanges.time[1],
+                phiMin=trackSelectorRanges.phi[0],
+                phiMax=trackSelectorRanges.phi[1],
+                etaMin=trackSelectorRanges.eta[0],
+                etaMax=trackSelectorRanges.eta[1],
+                absEtaMin=trackSelectorRanges.absEta[0],
+                absEtaMax=trackSelectorRanges.absEta[1],
+                ptMin=trackSelectorRanges.pt[0],
+                ptMax=trackSelectorRanges.pt[1],
+                removeCharged=trackSelectorRanges.removeCharged,
+                removeNeutral=trackSelectorRanges.removeNeutral,
+            ),
+        )
+        s.addAlgorithm(trackSelector)
+    else:
+        selectedTrackParameters = trackParameters
+
     inputParticles = "particles_input"
     outputVertices = "fittedVertices"
     selectedParticles = "particles_selected"
@@ -936,7 +989,7 @@ def addVertexFitting(
         fitVertices = VertexFitterAlgorithm(
             level=customLogLevel(),
             bField=field,
-            inputTrackParameters=trackParameters,
+            inputTrackParameters=selectedTrackParameters,
             inputProtoVertices=findVertices.config.outputProtoVertices,
             outputVertices=outputVertices,
         )
@@ -945,7 +998,7 @@ def addVertexFitting(
         findVertices = IterativeVertexFinderAlgorithm(
             level=customLogLevel(),
             bField=field,
-            inputTrackParameters=trackParameters,
+            inputTrackParameters=selectedTrackParameters,
             outputProtoVertices="protovertices",
             outputVertices=outputVertices,
         )
@@ -955,7 +1008,7 @@ def addVertexFitting(
         findVertices = AdaptiveMultiVertexFinderAlgorithm(
             level=customLogLevel(),
             bField=field,
-            inputTrackParameters=trackParameters,
+            inputTrackParameters=selectedTrackParameters,
             outputProtoVertices="protovertices",
             outputVertices=outputVertices,
             outputTime=outputTime,
@@ -978,8 +1031,8 @@ def addVertexFitting(
                 level=customLogLevel(),
                 inputAllTruthParticles=inputParticles,
                 inputSelectedTruthParticles=selectedParticles,
-                inputFittedTracks=trackParameters,
-                inputFittedTracksIndices="outputTrackIndices",
+                inputFittedTracks=selectedTrackParameters,
+                inputFittedTracksIndices=trackIndices,
                 inputAllFittedTracksTips="fittedTrackParametersTips",
                 inputMeasurementParticlesMap="measurement_particles_map",
                 inputTrajectories="trajectories" if trajectories is not None else "",

Examples/Scripts/Python/full_chain_itk.py

@@ -17,6 +17,7 @@
     CKFPerformanceConfig,
     addVertexFitting,
     VertexFinder,
+    TrackSelectorRanges,
 )
 
 ttbar_pu200 = False
@@ -92,21 +93,10 @@
     outputDirRoot=outputDir,
 )
 
-s.addAlgorithm(
-    acts.examples.TrackSelector(
-        level=acts.logging.INFO,
-        inputTrackParameters="fittedTrackParameters",
-        outputTrackParameters="trackparameters",
-        outputTrackIndices="outputTrackIndices",
-        removeNeutral=True,
-        absEtaMax=4.0,
-        ptMin=1.0 * u.GeV,
-    )
-)
-
 addVertexFitting(
     s,
     field,
+    TrackSelectorRanges(pt=(1.0 * u.GeV, None), absEta=(None, 4.0), removeNeutral=True),
     vertexFinder=VertexFinder.Iterative,
     outputDirRoot=outputDir,
     trajectories="trajectories",

Examples/Scripts/Python/full_chain_odd.py

@@ -17,6 +17,7 @@
     CKFPerformanceConfig,
     addVertexFitting,
     VertexFinder,
+    TrackSelectorRanges,
 )
 from common import getOpenDataDetectorDirectory
 from acts.examples.odd import getOpenDataDetector
@@ -100,21 +101,10 @@
     outputDirRoot=outputDir,
 )
 
-s.addAlgorithm(
-    acts.examples.TrackSelector(
-        level=acts.logging.INFO,
-        inputTrackParameters="fittedTrackParameters",
-        outputTrackParameters="trackparameters",
-        outputTrackIndices="outputTrackIndices",
-        removeNeutral=True,
-        absEtaMax=3.0,
-        ptMin=1.0 * u.GeV,
-    )
-)
-
 addVertexFitting(
     s,
     field,
+    TrackSelectorRanges(pt=(1.0 * u.GeV, None), absEta=(None, 3.0), removeNeutral=True),
     vertexFinder=VertexFinder.Iterative,
     outputDirRoot=outputDir,
     trajectories="trajectories",

Examples/Scripts/Python/vertex_fitting.py

@@ -5,7 +5,11 @@
 import acts
 from acts.examples import Sequencer, ParticleSelector, ParticleSmearing
 from acts.examples.simulation import addPythia8
-from acts.examples.reconstruction import addVertexFitting, VertexFinder
+from acts.examples.reconstruction import (
+    addVertexFitting,
+    VertexFinder,
+    TrackSelectorRanges,
+)
 
 u = acts.UnitConstants
 
@@ -53,7 +57,7 @@ def runVertexFitting(
     )
     s.addAlgorithm(ptclSelector)
 
-    trackParameters = "trackparameters"
+    trackParameters = "fittedTrackParameters"
     if inputTrackSummary is None or inputParticlePath is None:
         logger.info("Using smeared particles")
 
@@ -65,37 +69,34 @@ def runVertexFitting(
         )
         s.addAlgorithm(ptclSmearing)
         associatedParticles = selectedParticles
+
+        trackSelectorRanges = None
     else:
         logger.info("Reading track summary from %s", inputTrackSummary.resolve())
         assert inputTrackSummary.exists()
         associatedParticles = "associatedTruthParticles"
         trackSummaryReader = acts.examples.RootTrajectorySummaryReader(
             level=acts.logging.VERBOSE,
-            outputTracks="fittedTrackParameters",
+            outputTracks=trackParameters,
             outputParticles=associatedParticles,
             filePath=str(inputTrackSummary.resolve()),
             orderedEvents=False,
         )
         s.addReader(trackSummaryReader)
 
-        s.addAlgorithm(
-            acts.examples.TrackSelector(
-                level=acts.logging.INFO,
-                inputTrackParameters=trackSummaryReader.config.outputTracks,
-                outputTrackParameters=trackParameters,
-                outputTrackIndices="outputTrackIndices",
-                removeNeutral=True,
-                absEtaMax=2.5,
-                loc0Max=4.0 * u.mm,  # rho max
-                ptMin=500 * u.MeV,
-            )
+        trackSelectorRanges = TrackSelectorRanges(
+            removeNeutral=True,
+            absEta=(None, 2.5),
+            loc0=(None, 4.0 * u.mm),  # rho max
+            pt=(500 * u.MeV, None),
         )
 
     logger.info("Using vertex finder: %s", vertexFinder.name)
 
     addVertexFitting(
         s,
         field,
+        trackSelectorRanges=trackSelectorRanges,
         outputDirRoot=outputDir if outputRoot else None,
         associatedParticles=associatedParticles,
         trackParameters=trackParameters,