feat: add prototracks-to-seeds algorithm (#2219)

This allows to transform a prototracks collection into a seeds collection. Useful for the Exa.TrkX workflow.

Examples/Algorithms/Utilities/CMakeLists.txt

@@ -1,5 +1,6 @@
 add_library(
   ActsExamplesUtilities SHARED
+  src/PrototracksToSeeds.cpp
   src/SeedsToPrototracks.cpp
   src/TrajectoriesToPrototracks.cpp
   src/TracksToTrajectories.cpp)

Examples/Algorithms/Utilities/include/ActsExamples/Utilities/PrototracksToSeeds.hpp

@@ -0,0 +1,54 @@
+// This file is part of the Acts project.
+//
+// Copyright (C) 2023 CERN for the benefit of the Acts project
+//
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+#pragma once
+
+#include "ActsExamples/EventData/ProtoTrack.hpp"
+#include "ActsExamples/EventData/SimSeed.hpp"
+#include "ActsExamples/Framework/DataHandle.hpp"
+#include "ActsExamples/Framework/IAlgorithm.hpp"
+
+namespace ActsExamples {
+
+class PrototracksToSeeds final : public IAlgorithm {
+ public:
+  struct Config {
+    std::string inputProtoTracks;
+    std::string inputSpacePoints;
+    std::string outputSeeds = "seeds-from-prototracks";
+    std::string outputProtoTracks = "remaining-prototracks";
+  };
+
+  /// Construct the algorithm.
+  ///
+  /// @param cfg is the algorithm configuration
+  /// @param lvl is the logging level
+  PrototracksToSeeds(Config cfg, Acts::Logging::Level lvl);
+
+  /// Run the algorithm.
+  ///
+  /// @param ctx is the algorithm context with event information
+  /// @return a process code indication success or failure
+  ProcessCode execute(const AlgorithmContext& ctx) const final;
+
+  /// Const access to the config
+  const Config& config() const { return m_cfg; }
+
+ private:
+  Config m_cfg;
+
+  WriteDataHandle<SimSeedContainer> m_outputSeeds{this, "OutputSeeds"};
+  WriteDataHandle<ProtoTrackContainer> m_outputProtoTracks{this,
+                                                           "OutputProtoTracks"};
+  ReadDataHandle<SimSpacePointContainer> m_inputSpacePoints{this,
+                                                            "InputSpacePoints"};
+  ReadDataHandle<ProtoTrackContainer> m_inputProtoTracks{this,
+                                                         "InputProtoTracks"};
+};
+
+}  // namespace ActsExamples

Examples/Algorithms/Utilities/src/PrototracksToSeeds.cpp

@@ -0,0 +1,53 @@
+// This file is part of the Acts project.
+//
+// Copyright (C) 2023 CERN for the benefit of the Acts project
+//
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this
+// file, You can obtain one at http://mozilla.org/MPL/2.0/.
+
+#include "ActsExamples/Utilities/PrototracksToSeeds.hpp"
+
+#include "ActsExamples/EventData/IndexSourceLink.hpp"
+#include "ActsExamples/EventData/ProtoTrack.hpp"
+#include "ActsExamples/EventData/SimSeed.hpp"
+#include "ActsExamples/Framework/WhiteBoard.hpp"
+#include "ActsExamples/Utilities/EventDataTransforms.hpp"
+
+#include <algorithm>
+
+namespace ActsExamples {
+
+PrototracksToSeeds::PrototracksToSeeds(Config cfg, Acts::Logging::Level lvl)
+    : IAlgorithm("PrototracksToSeeds", lvl), m_cfg(std::move(cfg)) {
+  m_outputSeeds.initialize(m_cfg.outputSeeds);
+  m_outputProtoTracks.initialize(m_cfg.outputProtoTracks);
+  m_inputProtoTracks.initialize(m_cfg.inputProtoTracks);
+  m_inputSpacePoints.initialize(m_cfg.inputSpacePoints);
+}
+
+ProcessCode PrototracksToSeeds::execute(const AlgorithmContext& ctx) const {
+  auto prototracks = m_inputProtoTracks(ctx);
+
+  const auto nBefore = prototracks.size();
+  prototracks.erase(std::remove_if(prototracks.begin(), prototracks.end(),
+                                   [](const auto& t) { return t.size() < 3; }),
+                    prototracks.end());
+  ACTS_DEBUG("Discarded " << prototracks.size() - nBefore
+                          << " prototracks with less then 3 hits");
+
+  SimSeedContainer seeds;
+  seeds.reserve(prototracks.size());
+
+  const auto& sps = m_inputSpacePoints(ctx);
+  std::transform(prototracks.begin(), prototracks.end(),
+                 std::back_inserter(seeds),
+                 [&](const auto& pt) { return prototrackToSeed(pt, sps); });
+
+  m_outputSeeds(ctx, std::move(seeds));
+  m_outputProtoTracks(ctx, std::move(prototracks));
+
+  return ProcessCode::SUCCESS;
+}
+
+}  // namespace ActsExamples

Examples/Algorithms/Utilities/src/SeedsToPrototracks.cpp

@@ -18,15 +18,19 @@ namespace ActsExamples {
 struct AlgorithmContext;
 
 SeedsToPrototracks::SeedsToPrototracks(Config cfg, Acts::Logging::Level lvl)
-    : IAlgorithm("TrajectoriesToPrototracks", lvl), m_cfg(std::move(cfg)) {
+    : IAlgorithm("SeedsToPrototracks", lvl), m_cfg(std::move(cfg)) {
   m_inputSeeds.initialize(m_cfg.inputSeeds);
   m_outputProtoTracks.initialize(m_cfg.outputProtoTracks);
 }
 
 ProcessCode SeedsToPrototracks::execute(const AlgorithmContext& ctx) const {
   const auto seeds = m_inputSeeds(ctx);
 
-  auto tracks = seedsToPrototracks(seeds);
+  ProtoTrackContainer tracks;
+  tracks.reserve(seeds.size());
+
+  std::transform(seeds.begin(), seeds.end(), std::back_inserter(tracks),
+                 seedToPrototrack);
 
   m_outputProtoTracks(ctx, std::move(tracks));
 

Examples/Framework/include/ActsExamples/Utilities/EventDataTransforms.hpp

@@ -15,6 +15,10 @@ namespace ActsExamples {
 
 ProtoTrack seedToPrototrack(const SimSeed &seed);
 
-ProtoTrackContainer seedsToPrototracks(const SimSeedContainer &seeds);
+std::optional<SimSpacePoint> findSpacePointForIndex(
+    Index index, const SimSpacePointContainer &spacepoints);
+
+SimSeed prototrackToSeed(const ProtoTrack &track,
+                         const SimSpacePointContainer &spacepoints);
 
 }  // namespace ActsExamples

Examples/Framework/src/Utilities/EventDataTransforms.cpp

@@ -27,14 +27,56 @@ ActsExamples::ProtoTrack ActsExamples::seedToPrototrack(
   return track;
 }
 
-ActsExamples::ProtoTrackContainer ActsExamples::seedsToPrototracks(
-    const ActsExamples::SimSeedContainer& seeds) {
-  ProtoTrackContainer tracks;
-  tracks.reserve(seeds.size());
+std::optional<ActsExamples::SimSpacePoint> ActsExamples::findSpacePointForIndex(
+    ActsExamples::Index index, const SimSpacePointContainer& spacepoints) {
+  auto match = [&](const SimSpacePoint& sp) {
+    const auto& sls = sp.sourceLinks();
+    return std::any_of(sls.begin(), sls.end(), [&](const auto& sl) {
+      return sl.template get<IndexSourceLink>().index() == index;
+    });
+  };
 
-  for (const auto& seed : seeds) {
-    tracks.push_back(seedToPrototrack(seed));
+  auto found = std::find_if(spacepoints.begin(), spacepoints.end(), match);
+
+  if (found == spacepoints.end()) {
+    return std::nullopt;
+  }
+
+  return *found;
+}
+
+ActsExamples::SimSeed ActsExamples::prototrackToSeed(
+    const ActsExamples::ProtoTrack& track,
+    const ActsExamples::SimSpacePointContainer& spacepoints) {
+  auto findSpacePoint = [&](ActsExamples::Index index) {
+    auto found = findSpacePointForIndex(index, spacepoints);
+    if (not found) {
+      throw std::runtime_error("No spacepoint found for source-link index " +
+                               std::to_string(index));
+    }
+    return *found;
+  };
+
+  const auto s = track.size();
+  if (s < 3) {
+    throw std::runtime_error(
+        "Cannot convert track with less then 3 spacepoints to seed");
   }
 
-  return tracks;
+  std::vector<SimSpacePoint> ps;
+  ps.reserve(track.size());
+
+  std::transform(track.begin(), track.end(), std::back_inserter(ps),
+                 findSpacePoint);
+  std::sort(ps.begin(), ps.end(),
+            [](const auto& a, const auto& b) { return a.r() < b.r(); });
+
+  // Simply use r = m*z + t and solve for r=0 to find z vertex position...
+  // Probably not the textbook way to do
+  const auto m =
+      (ps.back().r() - ps.front().r()) / (ps.back().z() - ps.front().z());
+  const auto t = ps.front().r() - m * ps.front().z();
+  const auto z_vertex = -t / m;
+
+  return SimSeed(ps[0], ps[s / 2], ps[s - 1], z_vertex);
 }

Examples/Python/src/TrackFinding.cpp

@@ -24,6 +24,7 @@
 #include "ActsExamples/TrackFinding/SpacePointMaker.hpp"
 #include "ActsExamples/TrackFinding/TrackFindingAlgorithm.hpp"
 #include "ActsExamples/TrackFinding/TrackParamsEstimationAlgorithm.hpp"
+#include "ActsExamples/Utilities/PrototracksToSeeds.hpp"
 #include "ActsExamples/Utilities/SeedsToPrototracks.hpp"
 #include "ActsExamples/Utilities/TracksToTrajectories.hpp"
 #include "ActsExamples/Utilities/TrajectoriesToPrototracks.hpp"
@@ -341,6 +342,10 @@ void addTrackFinding(Context& ctx) {
   ACTS_PYTHON_DECLARE_ALGORITHM(ActsExamples::SeedsToPrototracks, mex,
                                 "SeedsToPrototracks", inputSeeds,
                                 outputProtoTracks);
+
+  ACTS_PYTHON_DECLARE_ALGORITHM(
+      ActsExamples::PrototracksToSeeds, mex, "PrototracksToSeeds",
+      inputProtoTracks, inputSpacePoints, outputSeeds, outputProtoTracks);
 }
 
 }  // namespace Acts::Python