feat: Add optional pass through of prototracks to param estimation al…

…gorithm (#2438)

In the case the param estimation fails, the number of parameters/seeds/prototracks does not match, which makes e.g., downstream track fitting impossible.
This allows to pass-through a prototrack collection through the algorithm, in order to ensure they match always.

Examples/Algorithms/TrackFinding/include/ActsExamples/TrackFinding/TrackParamsEstimationAlgorithm.hpp

@@ -52,8 +52,16 @@ class TrackParamsEstimationAlgorithm final : public IAlgorithm {
   struct Config {
     /// Input seeds collection.
     std::string inputSeeds;
+    /// Input prototracks (optional)
+    std::string inputProtoTracks;
     /// Output estimated track parameters collection.
     std::string outputTrackParameters;
+    /// Output seed collection - only seeds with successful parameter estimation
+    /// are propagated (optional)
+    std::string outputSeeds;
+    /// Output prototrack collection - only tracks with successful parameter
+    /// estimation are propagated (optional)
+    std::string outputProtoTracks;
     /// Tracking geometry for surface lookup.
     std::shared_ptr<const Acts::TrackingGeometry> trackingGeometry;
     /// Magnetic field variant.
@@ -99,9 +107,12 @@ class TrackParamsEstimationAlgorithm final : public IAlgorithm {
   Acts::BoundSquareMatrix m_covariance = Acts::BoundSquareMatrix::Zero();
 
   ReadDataHandle<SimSeedContainer> m_inputSeeds{this, "InputSeeds"};
+  ReadDataHandle<ProtoTrackContainer> m_inputTracks{this, "InputTracks"};
 
   WriteDataHandle<TrackParametersContainer> m_outputTrackParameters{
       this, "OutputTrackParameters"};
+  WriteDataHandle<SimSeedContainer> m_outputSeeds{this, "OutputSeeds"};
+  WriteDataHandle<ProtoTrackContainer> m_outputTracks{this, "OutputTracks"};
 };
 
 }  // namespace ActsExamples

Examples/Algorithms/TrackFinding/src/TrackParamsEstimationAlgorithm.cpp

@@ -50,7 +50,11 @@ ActsExamples::TrackParamsEstimationAlgorithm::TrackParamsEstimationAlgorithm(
   }
 
   m_inputSeeds.initialize(m_cfg.inputSeeds);
+  m_inputTracks.maybeInitialize(m_cfg.inputProtoTracks);
+
   m_outputTrackParameters.initialize(m_cfg.outputTrackParameters);
+  m_outputSeeds.maybeInitialize(m_cfg.outputSeeds);
+  m_outputTracks.maybeInitialize(m_cfg.outputProtoTracks);
 
   // Set up the track parameters covariance (the same for all tracks)
   m_covariance(Acts::eBoundLoc0, Acts::eBoundLoc0) =
@@ -81,6 +85,22 @@ ActsExamples::ProcessCode ActsExamples::TrackParamsEstimationAlgorithm::execute(
   TrackParametersContainer trackParameters;
   trackParameters.reserve(seeds.size());
 
+  std::optional<SimSeedContainer> outputSeeds;
+  if (m_outputSeeds.isInitialized()) {
+    outputSeeds->reserve(seeds.size());
+  }
+
+  const ProtoTrackContainer* inputTracks = nullptr;
+  std::optional<ProtoTrackContainer> outputTracks;
+  if (m_inputTracks.isInitialized() && m_outputTracks.isInitialized()) {
+    if (seeds.size() != inputTracks->size()) {
+      ACTS_FATAL("Inconsistent number of seeds and prototracks");
+      return ProcessCode::ABORT;
+    }
+    inputTracks = &m_inputTracks(ctx);
+    outputTracks->reserve(seeds.size());
+  }
+
   auto bCache = m_cfg.magneticField->makeCache(ctx.magFieldContext);
 
   IndexSourceLink::SurfaceAccessor surfaceAccessor{*m_cfg.trackingGeometry};
@@ -125,11 +145,25 @@ ActsExamples::ProcessCode ActsExamples::TrackParamsEstimationAlgorithm::execute(
       double charge = std::copysign(1, params[Acts::eBoundQOverP]);
       trackParameters.emplace_back(surface->getSharedPtr(), params, charge,
                                    m_covariance);
+      if (outputSeeds) {
+        outputSeeds->push_back(seed);
+      }
+      if (outputTracks && inputTracks != nullptr) {
+        outputTracks->push_back(inputTracks->at(iseed));
+      }
     }
   }
 
   ACTS_VERBOSE("Estimated " << trackParameters.size() << " track parameters");
 
   m_outputTrackParameters(ctx, std::move(trackParameters));
+  if (m_outputSeeds.isInitialized()) {
+    m_outputSeeds(ctx, std::move(*outputSeeds));
+  }
+
+  if (m_outputTracks.isInitialized()) {
+    m_outputTracks(ctx, std::move(*outputTracks));
+  }
+
   return ProcessCode::SUCCESS;
 }

Examples/Python/src/TrackFinding.cpp

@@ -259,9 +259,10 @@ void addTrackFinding(Context& ctx) {
 
   ACTS_PYTHON_DECLARE_ALGORITHM(
       ActsExamples::TrackParamsEstimationAlgorithm, mex,
-      "TrackParamsEstimationAlgorithm", inputSeeds, outputTrackParameters,
-      trackingGeometry, magneticField, bFieldMin, sigmaLoc0, sigmaLoc1,
-      sigmaPhi, sigmaTheta, sigmaQOverP, sigmaT0, initialVarInflation);
+      "TrackParamsEstimationAlgorithm", inputSeeds, inputProtoTracks,
+      outputTrackParameters, outputSeeds, outputProtoTracks, trackingGeometry,
+      magneticField, bFieldMin, sigmaLoc0, sigmaLoc1, sigmaPhi, sigmaTheta,
+      sigmaQOverP, sigmaT0, initialVarInflation);
 
   {
     using Alg = ActsExamples::TrackFindingAlgorithm;