refactor!: Add ability to copy Tracks between containers (#1980)

This preserves the dynamic columns if the backend supports this. Also adds a method to the `TrackContainer` layer that should ensure the backend copies the types using the correct types.

This PR also switches the `TrackSelector` algorithm to use this, instead of copying the input tracks and removing the discarded tracks. In my testing, this brings the selector performance back in line with expectations on ttbar.

BREAKING CHANGE: `TrackContainer` backends needs to implement these methods:
- `copyDynamicFrom_impl`
- `ensureDynamicColumns_impl`

Core/include/Acts/EventData/Track.hpp

@@ -330,6 +330,29 @@ class TrackProxy {
     return *m_container;
   }
 
+  /// Copy the content of another track proxy into this one
+  /// @tparam track_proxy_t the other track proxy's type
+  /// @param other The the track proxy
+  template <typename track_proxy_t, bool RO = ReadOnly,
+            typename = std::enable_if_t<!RO>>
+  void copyFrom(const track_proxy_t& other) {
+    // @TODO: Add constraint on which track proxies are allowed,
+    // this is only implicit right now
+
+    tipIndex() = other.tipIndex();
+    parameters() = other.parameters();
+    covariance() = other.covariance();
+    if (other.hasReferenceSurface()) {
+      setReferenceSurface(other.referenceSurface().getSharedPtr());
+    }
+    nMeasurements() = other.nMeasurements();
+    nHoles() = other.nHoles();
+
+    // This will only be valid if the backends match and support this operation
+    m_container->copyDynamicFrom(m_index, other.m_container->container(),
+                                 other.m_index);
+  }
+
   /// Return a reference to the track container backend, const version.
   /// @return reference to the track container backend
   const auto& container() const { return *m_container; }
@@ -341,6 +364,10 @@ class TrackProxy {
     return &(*m_container) == &(*other.m_container) && m_index == other.m_index;
   }
 
+  // Track proxies are friends, not food!
+  template <typename A, typename B, template <typename> class H, bool R>
+  friend class TrackProxy;
+
  private:
   TrackProxy(ConstIf<TrackContainer<Container, Trajectory, holder_t>, ReadOnly>&
                  container,
@@ -635,6 +662,18 @@ class TrackContainer {
                                          ConstTrackProxy, true>{*this, size()};
   }
 
+  /// Helper function to make this track container match the dynamic columns of
+  /// another one. This will only work if the track container supports this
+  /// source, and depends on the implementation details of the dynamic columns
+  /// of the container
+  /// @tparam other_track_container_t Type of the other track container
+  /// @param other The other track container
+  template <typename other_track_container_t, bool RO = ReadOnly,
+            typename = std::enable_if_t<!RO>>
+  void ensureDynamicColumns(const other_track_container_t& other) {
+    container().ensureDynamicColumns_impl(other.container());
+  }
+
  protected:
   template <typename T, HashedString key, bool RO = ReadOnly,
             typename = std::enable_if_t<!RO>>
@@ -689,6 +728,11 @@ class TrackContainer {
   }
 
  private:
+  template <typename T, bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
+  void copyDynamicFrom(IndexType dstIdx, const T& src, IndexType srcIdx) {
+    container().copyDynamicFrom_impl(dstIdx, src, srcIdx);
+  }
+
   detail_tc::ConstIf<holder_t<track_container_t>, ReadOnly> m_container;
   detail_tc::ConstIf<holder_t<traj_t>, ReadOnly> m_traj;
 };

Core/include/Acts/EventData/VectorTrackContainer.hpp

@@ -187,6 +187,15 @@ class VectorTrackContainer final : public detail_vtc::VectorTrackContainerBase {
     return ConstCovariance{m_cov[itrack].data()};
   }
 
+  void copyDynamicFrom_impl(IndexType dstIdx,
+                            const VectorTrackContainerBase& src,
+                            IndexType srcIdx);
+
+  void ensureDynamicColumns_impl(
+      const detail_vtc::VectorTrackContainerBase& other);
+
+  void reserve(IndexType size);
+
   // END INTERFACE
 };
 

Core/include/Acts/EventData/detail/DynamicColumn.hpp

@@ -21,10 +21,14 @@ struct DynamicColumnBase {
 
   virtual void add() = 0;
   virtual void clear() = 0;
+  virtual void reserve(size_t size) = 0;
   virtual void erase(size_t i) = 0;
   virtual size_t size() const = 0;
+  virtual void copyFrom(size_t dstIdx, const DynamicColumnBase& src,
+                        size_t srcIdx) = 0;
 
-  virtual std::unique_ptr<DynamicColumnBase> clone() const = 0;
+  virtual std::unique_ptr<DynamicColumnBase> clone(
+      bool empty = false) const = 0;
 };
 
 template <typename T>
@@ -41,13 +45,25 @@ struct DynamicColumn : public DynamicColumnBase {
 
   void add() override { m_vector.emplace_back(); }
   void clear() override { m_vector.clear(); }
+  void reserve(size_t size) override { m_vector.reserve(size); }
   void erase(size_t i) override { m_vector.erase(m_vector.begin() + i); }
   size_t size() const override { return m_vector.size(); }
 
-  std::unique_ptr<DynamicColumnBase> clone() const override {
+  std::unique_ptr<DynamicColumnBase> clone(bool empty) const override {
+    if (empty) {
+      return std::make_unique<DynamicColumn<T>>();
+    }
     return std::make_unique<DynamicColumn<T>>(*this);
   }
 
+  void copyFrom(size_t dstIdx, const DynamicColumnBase& src,
+                size_t srcIdx) override {
+    const auto* other = dynamic_cast<const DynamicColumn<T>*>(&src);
+    assert(other != nullptr &&
+           "Source column is not of same type as destination");
+    m_vector.at(dstIdx) = other->m_vector.at(srcIdx);
+  }
+
   std::vector<T> m_vector;
 };
 
@@ -68,14 +84,26 @@ struct DynamicColumn<bool> : public DynamicColumnBase {
   }
 
   void add() override { m_vector.emplace_back(); }
+  void reserve(size_t size) override { m_vector.reserve(size); }
   void clear() override { m_vector.clear(); }
   void erase(size_t i) override { m_vector.erase(m_vector.begin() + i); }
   size_t size() const override { return m_vector.size(); }
 
-  std::unique_ptr<DynamicColumnBase> clone() const override {
+  std::unique_ptr<DynamicColumnBase> clone(bool empty) const override {
+    if (empty) {
+      return std::make_unique<DynamicColumn<bool>>();
+    }
     return std::make_unique<DynamicColumn<bool>>(*this);
   }
 
+  void copyFrom(size_t dstIdx, const DynamicColumnBase& src,
+                size_t srcIdx) override {
+    const auto* other = dynamic_cast<const DynamicColumn<bool>*>(&src);
+    assert(other != nullptr &&
+           "Source column is not of same type as destination");
+    m_vector.at(dstIdx) = other->m_vector.at(srcIdx);
+  }
+
   std::vector<Wrapper> m_vector;
 };
 

Core/src/EventData/VectorTrackContainer.cpp

@@ -71,4 +71,41 @@ void VectorTrackContainer::removeTrack_impl(IndexType itrack) {
   }
 }
 
+void VectorTrackContainer::copyDynamicFrom_impl(
+    IndexType dstIdx, const VectorTrackContainerBase& src, IndexType srcIdx) {
+  for (const auto& [key, value] : src.m_dynamic) {
+    auto it = m_dynamic.find(key);
+    if (it == m_dynamic.end()) {
+      throw std::invalid_argument{
+          "Destination container does not have matching dynamic column"};
+    }
+
+    it->second->copyFrom(dstIdx, *value, srcIdx);
+  }
+}
+
+void VectorTrackContainer::ensureDynamicColumns_impl(
+    const detail_vtc::VectorTrackContainerBase& other) {
+  for (auto& [key, value] : other.m_dynamic) {
+    if (m_dynamic.find(key) == m_dynamic.end()) {
+      m_dynamic[key] = value->clone(true);
+    }
+  }
+}
+
+void VectorTrackContainer::reserve(IndexType size) {
+  m_tipIndex.reserve(size);
+
+  m_params.reserve(size);
+  m_cov.reserve(size);
+  m_referenceSurfaces.reserve(size);
+
+  m_nMeasurements.reserve(size);
+  m_nHoles.reserve(size);
+
+  for (auto& [key, vec] : m_dynamic) {
+    vec->reserve(size);
+  }
+}
+
 }  // namespace Acts

Examples/Algorithms/TruthTracking/ActsExamples/TruthTracking/TrackSelector.cpp

@@ -60,37 +60,26 @@ ActsExamples::ProcessCode ActsExamples::TrackSelector::execute(
   std::shared_ptr<Acts::ConstVectorMultiTrajectory> trackStateContainer =
       inputTracks.trackStateContainerHolder();
 
-  auto trackContainer = std::make_shared<Acts::VectorTrackContainer>(
-      inputTracks.container());  // mutable copy of the immutable
-                                 // source track container
+  auto trackContainer = std::make_shared<Acts::VectorTrackContainer>();
+
+  // temporary empty track state container: we don't change the original one,
+  // but we need one for filtering
   auto tempTrackStateContainer =
       std::make_shared<Acts::VectorMultiTrajectory>();
 
   TrackContainer filteredTracks{trackContainer, tempTrackStateContainer};
+  filteredTracks.ensureDynamicColumns(inputTracks);
 
-  ACTS_VERBOSE(
-      "Track container size before filtering: " << filteredTracks.size());
+  trackContainer->reserve(inputTracks.size());
 
-  size_t nRemoved = 0;
+  ACTS_VERBOSE("Track container size before filtering: " << inputTracks.size());
 
-  for (Acts::MultiTrajectoryTraits::IndexType itrack = 0;
-       itrack < filteredTracks.size();) {
-    auto track = filteredTracks.getTrack(itrack);
-    if (isValidTrack(track)) {
-      // Track is valid, go to next index
-      ACTS_VERBOSE(" - Keeping track #" << itrack);
-      itrack++;
+  for (auto track : inputTracks) {
+    if (!isValidTrack(track)) {
       continue;
     }
-    ACTS_VERBOSE(" - Removing track #" << itrack);
-    filteredTracks.removeTrack(itrack);
-    nRemoved++;
-    // Do not increment track index
-  }
-
-  if (nRemoved != inputTracks.size() - filteredTracks.size()) {
-    ACTS_ERROR("Inconsistent track container size after removing "
-               << nRemoved << " tracks");
+    auto destProxy = filteredTracks.getTrack(filteredTracks.addTrack());
+    destProxy.copyFrom(track);
   }
 
   ACTS_VERBOSE(

Tests/UnitTests/Core/EventData/TrackTests.cpp

@@ -443,4 +443,79 @@ BOOST_AUTO_TEST_CASE_TEMPLATE(DynamicColumns, factory_t, holder_types) {
   BOOST_CHECK_EQUAL((t.template component<float, "col_a"_hash>()), 5.6f);
 }
 
+BOOST_AUTO_TEST_CASE(CopyTracksIncludingDynamicColumns) {
+  // mutable source
+  VectorTrackContainer vtc{};
+  VectorMultiTrajectory mtj{};
+  TrackContainer tc{vtc, mtj};
+  tc.addColumn<size_t>("counter");
+  tc.addColumn<bool>("odd");
+
+  TrackContainer tc2{VectorTrackContainer{}, VectorMultiTrajectory{}};
+  // doesn't have the dynamic column
+
+  TrackContainer tc3{VectorTrackContainer{}, VectorMultiTrajectory{}};
+  tc3.addColumn<size_t>("counter");
+  tc3.addColumn<bool>("odd");
+
+  for (size_t i = 0; i < 10; i++) {
+    auto t = tc.getTrack(tc.addTrack());
+    t.tipIndex() = i;
+    t.template component<size_t>("counter") = i;
+    t.template component<bool>("odd") = i % 2 == 0;
+
+    auto t2 = tc2.getTrack(tc2.addTrack());
+    BOOST_CHECK_THROW(t2.copyFrom(t),
+                      std::invalid_argument);  // this should fail
+
+    auto t3 = tc3.getTrack(tc3.addTrack());
+    t3.copyFrom(t);  // this should work
+
+    BOOST_CHECK_EQUAL(t.tipIndex(), t3.tipIndex());
+    BOOST_CHECK_EQUAL(t.template component<size_t>("counter"),
+                      t3.template component<size_t>("counter"));
+    BOOST_CHECK_EQUAL(t.template component<bool>("odd"),
+                      t3.template component<bool>("odd"));
+  }
+
+  TrackContainer tc4{ConstVectorTrackContainer{vtc},
+                     ConstVectorMultiTrajectory{}};
+
+  TrackContainer tc5{VectorTrackContainer{}, VectorMultiTrajectory{}};
+  tc5.addColumn<size_t>("counter");
+  tc5.addColumn<bool>("odd");
+
+  for (size_t i = 0; i < 10; i++) {
+    auto t4 = tc4.getTrack(i);  // const source!
+
+    auto t5 = tc5.getTrack(tc5.addTrack());
+    t5.copyFrom(t4);  // this should work
+
+    BOOST_CHECK_EQUAL(t4.tipIndex(), t5.tipIndex());
+    BOOST_CHECK_EQUAL(t4.template component<size_t>("counter"),
+                      t5.template component<size_t>("counter"));
+    BOOST_CHECK_EQUAL(t4.template component<bool>("odd"),
+                      t5.template component<bool>("odd"));
+  }
+}
+
+BOOST_AUTO_TEST_CASE(EnsureDynamicColumns) {
+  TrackContainer tc{VectorTrackContainer{}, VectorMultiTrajectory{}};
+  tc.addColumn<size_t>("counter");
+  tc.addColumn<bool>("odd");
+
+  BOOST_CHECK(tc.hasColumn("counter"));
+  BOOST_CHECK(tc.hasColumn("odd"));
+
+  TrackContainer tc2{VectorTrackContainer{}, VectorMultiTrajectory{}};
+
+  BOOST_CHECK(!tc2.hasColumn("counter"));
+  BOOST_CHECK(!tc2.hasColumn("odd"));
+
+  tc2.ensureDynamicColumns(tc);
+
+  BOOST_CHECK(tc2.hasColumn("counter"));
+  BOOST_CHECK(tc2.hasColumn("odd"));
+}
+
 BOOST_AUTO_TEST_SUITE_END()