feat: Method to add components to an existing track state

Core/include/Acts/EventData/MultiTrajectory.hpp

@@ -556,6 +556,15 @@ class MultiTrajectory {
     self().unset_impl(target, istate);
   }
 
+  /// Add additional components to an existing track state
+  /// @note Only available if the track state container is not read-only
+  /// @param istate The track state index to alter
+  /// @param mask The bitmask that instructs which components to allocate
+  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
+  void addTrackStateComponents(IndexType istate, TrackStatePropMask mask) {
+    self().addTrackStateComponents_impl(istate, mask);
+  }
+
   /// Retrieve a mutable reference to a component
   /// @tparam T The type of the component to access
   /// @tparam key String key for the component to access

Core/include/Acts/EventData/TrackStateProxy.hpp

@@ -298,6 +298,14 @@ class TrackStateProxy {
     m_traj->self().unset(target, m_istate);
   }
 
+  /// Add additional components to the track state
+  /// @note Only available if the track state proxy is not read-only
+  /// @param mask The bitmask that instructs which components to allocate
+  template <bool RO = ReadOnly, typename = std::enable_if_t<!RO>>
+  void addComponents(TrackStatePropMask mask) {
+    m_traj->self().addTrackStateComponents_impl(m_istate, mask);
+  }
+
   /// Reference surface.
   /// @return the reference surface
   const Surface& referenceSurface() const {

Core/include/Acts/EventData/VectorMultiTrajectory.hpp

@@ -424,6 +424,8 @@ class VectorMultiTrajectory final
       TrackStatePropMask mask = TrackStatePropMask::All,
       IndexType iprevious = kInvalid);
 
+  void addTrackStateComponents_impl(IndexType istate, TrackStatePropMask mask);
+
   void reserve(std::size_t n);
 
   void shareFrom_impl(IndexType iself, IndexType iother,

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

@@ -254,6 +254,69 @@ class MultiTrajectoryTestsCommon {
     alwaysPresent(ts);
   }
 
+  void testAddTrackStateComponents() {
+    using PM = TrackStatePropMask;
+
+    trajectory_t t = m_factory.create();
+
+    auto ts = t.makeTrackState(PM::None);
+    BOOST_CHECK(!ts.hasPredicted());
+    BOOST_CHECK(!ts.hasFiltered());
+    BOOST_CHECK(!ts.hasSmoothed());
+    BOOST_CHECK(!ts.hasCalibrated());
+    BOOST_CHECK(!ts.hasJacobian());
+
+    ts.addComponents(PM::None);
+    BOOST_CHECK(!ts.hasPredicted());
+    BOOST_CHECK(!ts.hasFiltered());
+    BOOST_CHECK(!ts.hasSmoothed());
+    BOOST_CHECK(!ts.hasCalibrated());
+    BOOST_CHECK(!ts.hasJacobian());
+
+    ts.addComponents(PM::Predicted);
+    BOOST_CHECK(ts.hasPredicted());
+    BOOST_CHECK(!ts.hasFiltered());
+    BOOST_CHECK(!ts.hasSmoothed());
+    BOOST_CHECK(!ts.hasCalibrated());
+    BOOST_CHECK(!ts.hasJacobian());
+
+    ts.addComponents(PM::Filtered);
+    BOOST_CHECK(ts.hasPredicted());
+    BOOST_CHECK(ts.hasFiltered());
+    BOOST_CHECK(!ts.hasSmoothed());
+    BOOST_CHECK(!ts.hasCalibrated());
+    BOOST_CHECK(!ts.hasJacobian());
+
+    ts.addComponents(PM::Smoothed);
+    BOOST_CHECK(ts.hasPredicted());
+    BOOST_CHECK(ts.hasFiltered());
+    BOOST_CHECK(ts.hasSmoothed());
+    BOOST_CHECK(!ts.hasCalibrated());
+    BOOST_CHECK(!ts.hasJacobian());
+
+    ts.addComponents(PM::Calibrated);
+    ts.allocateCalibrated(5);
+    BOOST_CHECK(ts.hasPredicted());
+    BOOST_CHECK(ts.hasFiltered());
+    BOOST_CHECK(ts.hasSmoothed());
+    BOOST_CHECK(ts.hasCalibrated());
+    BOOST_CHECK(!ts.hasJacobian());
+
+    ts.addComponents(PM::Jacobian);
+    BOOST_CHECK(ts.hasPredicted());
+    BOOST_CHECK(ts.hasFiltered());
+    BOOST_CHECK(ts.hasSmoothed());
+    BOOST_CHECK(ts.hasCalibrated());
+    BOOST_CHECK(ts.hasJacobian());
+
+    ts.addComponents(PM::All);
+    BOOST_CHECK(ts.hasPredicted());
+    BOOST_CHECK(ts.hasFiltered());
+    BOOST_CHECK(ts.hasSmoothed());
+    BOOST_CHECK(ts.hasCalibrated());
+    BOOST_CHECK(ts.hasJacobian());
+  }
+
   void testTrackStateProxyCrossTalk(std::default_random_engine& rng) {
     TestTrackState pc(rng, 2u);
 

Core/src/EventData/VectorMultiTrajectory.cpp

@@ -87,6 +87,56 @@ auto VectorMultiTrajectory::addTrackState_impl(TrackStatePropMask mask,
   return index;
 }
 
+void VectorMultiTrajectory::addTrackStateComponents_impl(
+    IndexType istate, TrackStatePropMask mask) {
+  using PropMask = TrackStatePropMask;
+
+  IndexData& p = m_index[istate];
+  PropMask currentMask = p.allocMask;
+
+  assert(m_params.size() == m_cov.size());
+
+  if (ACTS_CHECK_BIT(mask, PropMask::Predicted) &&
+      !ACTS_CHECK_BIT(currentMask, PropMask::Predicted)) {
+    m_params.emplace_back();
+    m_cov.emplace_back();
+    p.ipredicted = m_params.size() - 1;
+  }
+
+  if (ACTS_CHECK_BIT(mask, PropMask::Filtered) &&
+      !ACTS_CHECK_BIT(currentMask, PropMask::Filtered)) {
+    m_params.emplace_back();
+    m_cov.emplace_back();
+    p.ifiltered = m_params.size() - 1;
+  }
+
+  if (ACTS_CHECK_BIT(mask, PropMask::Smoothed) &&
+      !ACTS_CHECK_BIT(currentMask, PropMask::Smoothed)) {
+    m_params.emplace_back();
+    m_cov.emplace_back();
+    p.ismoothed = m_params.size() - 1;
+  }
+
+  assert(m_params.size() == m_cov.size());
+
+  if (ACTS_CHECK_BIT(mask, PropMask::Jacobian) &&
+      !ACTS_CHECK_BIT(currentMask, PropMask::Jacobian)) {
+    m_jac.emplace_back();
+    p.ijacobian = m_jac.size() - 1;
+  }
+
+  if (ACTS_CHECK_BIT(mask, PropMask::Calibrated) &&
+      !ACTS_CHECK_BIT(currentMask, PropMask::Calibrated)) {
+    m_sourceLinks.emplace_back(std::nullopt);
+    p.icalibratedsourcelink = m_sourceLinks.size() - 1;
+
+    m_projectors.emplace_back();
+    p.iprojector = m_projectors.size() - 1;
+  }
+
+  p.allocMask |= mask;
+}
+
 void VectorMultiTrajectory::shareFrom_impl(IndexType iself, IndexType iother,
                                            TrackStatePropMask shareSource,
                                            TrackStatePropMask shareTarget) {

Plugins/Podio/include/Acts/Plugins/Podio/PodioTrackStateContainer.hpp

@@ -484,6 +484,39 @@ class MutablePodioTrackStateContainer final
     return m_collection->size() - 1;
   }
 
+  void addTrackStateComponents_impl(IndexType istate, TrackStatePropMask mask) {
+    auto& data = m_collection->at(istate).data();
+
+    if (ACTS_CHECK_BIT(mask, TrackStatePropMask::Predicted) &&
+        data.ipredicted == kInvalid) {
+      m_params->create();
+      data.ipredicted = m_params->size() - 1;
+    }
+
+    if (ACTS_CHECK_BIT(mask, TrackStatePropMask::Filtered) &&
+        data.ifiltered == kInvalid) {
+      m_params->create();
+      data.ifiltered = m_params->size() - 1;
+    }
+
+    if (ACTS_CHECK_BIT(mask, TrackStatePropMask::Smoothed) &&
+        data.ismoothed == kInvalid) {
+      m_params->create();
+      data.ismoothed = m_params->size() - 1;
+    }
+
+    if (ACTS_CHECK_BIT(mask, TrackStatePropMask::Jacobian) &&
+        data.ijacobian == kInvalid) {
+      m_jacs->create();
+      data.ijacobian = m_jacs->size() - 1;
+    }
+
+    if (ACTS_CHECK_BIT(mask, TrackStatePropMask::Calibrated) &&
+        !data.hasProjector) {
+      data.hasProjector = true;
+    }
+  }
+
   void shareFrom_impl(TrackIndexType iself, TrackIndexType iother,
                       TrackStatePropMask shareSource,
                       TrackStatePropMask shareTarget) {

Tests/UnitTests/Core/EventData/MultiTrajectoryTests.cpp

@@ -131,6 +131,11 @@ BOOST_AUTO_TEST_CASE(AddTrackStateWithBitMask) {
   ct.testAddTrackStateWithBitMask();
 }
 
+BOOST_AUTO_TEST_CASE(AddTrackStateComponents) {
+  CommonTests ct;
+  ct.testAddTrackStateComponents();
+}
+
 // assert expected "cross-talk" between trackstate proxies
 BOOST_AUTO_TEST_CASE(TrackStateProxyCrossTalk) {
   CommonTests ct;

Tests/UnitTests/Plugins/Podio/PodioTrackStateContainerTest.cpp

@@ -132,6 +132,11 @@ BOOST_AUTO_TEST_CASE(AddTrackStateWithBitMask) {
   ct.testAddTrackStateWithBitMask();
 }
 
+BOOST_AUTO_TEST_CASE(AddTrackStateComponents) {
+  CommonTests ct;
+  ct.testAddTrackStateComponents();
+}
+
 // assert expected "cross-talk" between trackstate proxies
 BOOST_AUTO_TEST_CASE(TrackStateProxyCrossTalk) {
   CommonTests ct;