feat: track modifier for examples (#1726)

while dealing with vertexing issues I thought adding a track parameter modification algorithm to the examples framework might be useful

another try of #1658

Core/include/Acts/EventData/SingleBoundTrackParameters.hpp

@@ -151,9 +151,13 @@ class SingleBoundTrackParameters {
       default;
   SingleBoundTrackParameters& operator=(SingleBoundTrackParameters&&) = default;
 
+  /// Parameters vector.
+  ParametersVector& parameters() { return m_params; }
   /// Parameters vector.
   const ParametersVector& parameters() const { return m_params; }
   /// Optional covariance matrix.
+  std::optional<CovarianceMatrix>& covariance() { return m_cov; }
+  /// Optional covariance matrix.
   const std::optional<CovarianceMatrix>& covariance() const { return m_cov; }
 
   /// Access a single parameter value indentified by its index.

Examples/Algorithms/TruthTracking/ActsExamples/TruthTracking/TrackModifier.cpp

@@ -0,0 +1,116 @@
+// This file is part of the Acts project.
+//
+// Copyright (C) 2022 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/TruthTracking/TrackModifier.hpp"
+
+#include "Acts/EventData/MultiTrajectory.hpp"
+#include "Acts/Utilities/ThrowAssert.hpp"
+#include "ActsExamples/EventData/Track.hpp"
+#include "ActsExamples/EventData/Trajectories.hpp"
+#include "ActsExamples/Framework/WhiteBoard.hpp"
+
+#include <cmath>
+#include <cstdint>
+#include <stdexcept>
+#include <vector>
+
+ActsExamples::TrackModifier::TrackModifier(const Config& config,
+                                           Acts::Logging::Level level)
+    : BareAlgorithm("TrackModifier", level), m_cfg(config) {
+  if (m_cfg.inputTrajectories.empty() == m_cfg.inputTrackParameters.empty()) {
+    throw std::invalid_argument(
+        "Exactly one of trajectories or track parameters input must be set");
+  }
+  if (m_cfg.outputTrajectories.empty() == m_cfg.outputTrackParameters.empty()) {
+    throw std::invalid_argument(
+        "Exactly one of trajectories or track parameters output must be set");
+  }
+  if (m_cfg.inputTrajectories.empty() != m_cfg.outputTrajectories.empty()) {
+    throw std::invalid_argument(
+        "Input and output for trajectories and track parameters have to be "
+        "used consistently");
+  }
+}
+
+ActsExamples::ProcessCode ActsExamples::TrackModifier::execute(
+    const ActsExamples::AlgorithmContext& ctx) const {
+  auto modifyTrack = [this](auto trk) {
+    {
+      auto& params = trk.parameters();
+
+      if (m_cfg.killTime) {
+        params[Acts::eBoundTime] = 0;
+      }
+    }
+
+    {
+      auto& optCov = trk.covariance();
+
+      if (optCov) {
+        auto& cov = *optCov;
+
+        if (m_cfg.dropCovariance) {
+          cov = Acts::BoundSymMatrix(cov.diagonal().asDiagonal());
+        }
+        if (m_cfg.covScale != 1) {
+          cov *= m_cfg.covScale;
+        }
+        if (m_cfg.killTime) {
+          cov.row(Acts::eBoundTime).setZero();
+          cov.col(Acts::eBoundTime).setZero();
+          cov(Acts::eBoundTime, Acts::eBoundTime) = 1;
+        }
+      }
+    }
+
+    return trk;
+  };
+
+  if (!m_cfg.inputTrackParameters.empty()) {
+    const auto& inputTrackParameters =
+        ctx.eventStore.get<TrackParametersContainer>(
+            m_cfg.inputTrackParameters);
+    TrackParametersContainer outputTrackParameters;
+    outputTrackParameters.reserve(inputTrackParameters.size());
+
+    for (uint32_t i = 0; i < inputTrackParameters.size(); ++i) {
+      const auto& trk = inputTrackParameters[i];
+      outputTrackParameters.push_back(modifyTrack(trk));
+    }
+
+    ctx.eventStore.add(m_cfg.outputTrackParameters,
+                       std::move(outputTrackParameters));
+  } else if (!m_cfg.inputTrajectories.empty()) {
+    const auto& inputTrajectories =
+        ctx.eventStore.get<TrajectoriesContainer>(m_cfg.inputTrajectories);
+    TrajectoriesContainer outputTrajectories;
+    outputTrajectories.reserve(inputTrajectories.size());
+
+    for (const auto& trajectories : inputTrajectories) {
+      std::vector<Acts::MultiTrajectoryTraits::IndexType> tips;
+      tips.reserve(trajectories.tips().size());
+      Trajectories::IndexedParameters parameters;
+
+      for (auto tip : trajectories.tips()) {
+        if (!trajectories.hasTrackParameters(tip)) {
+          continue;
+        }
+        const auto& trk = trajectories.trackParameters(tip);
+        tips.push_back(tip);
+        parameters.emplace(tip, modifyTrack(trk));
+      }
+
+      outputTrajectories.emplace_back(trajectories.multiTrajectoryPtr(), tips,
+                                      parameters);
+    }
+
+    ctx.eventStore.add(m_cfg.outputTrajectories, std::move(outputTrajectories));
+  }
+
+  return ProcessCode::SUCCESS;
+}

Examples/Algorithms/TruthTracking/ActsExamples/TruthTracking/TrackModifier.hpp

@@ -0,0 +1,54 @@
+// This file is part of the Acts project.
+//
+// Copyright (C) 2022 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/Framework/BareAlgorithm.hpp"
+
+#include <limits>
+#include <string>
+
+namespace ActsExamples {
+
+/// Select tracks by applying some selection cuts.
+class TrackModifier final : public BareAlgorithm {
+ public:
+  struct Config {
+    /// Optional. Input trajectories container. Mutually exclusive with track
+    /// parameters input.
+    std::string inputTrajectories;
+    /// Optional. Input track parameters collection. Mutually exclusive with
+    /// trajectories input.
+    std::string inputTrackParameters;
+    /// Optional. Output trajectories container. Will only be set if
+    /// trajectories input was set
+    std::string outputTrajectories;
+    /// Optional. Output track parameters collection. Will only be set if track
+    /// parameters input was set.
+    std::string outputTrackParameters;
+
+    /// When turned on, only keed the diagonal of the cov matrix.
+    bool dropCovariance{false};
+    /// Scale cov matrix;
+    double covScale{1};
+    /// Remove all time components
+    bool killTime{false};
+  };
+
+  TrackModifier(const Config& config, Acts::Logging::Level level);
+
+  ProcessCode execute(const AlgorithmContext& ctx) const final;
+
+  /// Get readonly access to the config parameters
+  const Config& config() const { return m_cfg; }
+
+ private:
+  Config m_cfg;
+};
+
+}  // namespace ActsExamples

Examples/Algorithms/TruthTracking/CMakeLists.txt

@@ -3,6 +3,7 @@ add_library(
   ActsExamples/TruthTracking/ParticleSelector.cpp
   ActsExamples/TruthTracking/ParticleSmearing.cpp
   ActsExamples/TruthTracking/TrackSelector.cpp
+  ActsExamples/TruthTracking/TrackModifier.cpp
   ActsExamples/TruthTracking/TruthSeedSelector.cpp
   ActsExamples/TruthTracking/TruthTrackFinder.cpp
   ActsExamples/TruthTracking/TruthVertexFinder.cpp

Examples/Python/src/TruthTracking.cpp

@@ -9,6 +9,7 @@
 #include "Acts/Plugins/Python/Utilities.hpp"
 #include "ActsExamples/TruthTracking/ParticleSelector.hpp"
 #include "ActsExamples/TruthTracking/ParticleSmearing.hpp"
+#include "ActsExamples/TruthTracking/TrackModifier.hpp"
 #include "ActsExamples/TruthTracking/TrackSelector.hpp"
 #include "ActsExamples/TruthTracking/TruthSeedSelector.hpp"
 #include "ActsExamples/TruthTracking/TruthTrackFinder.hpp"
@@ -170,6 +171,11 @@ void addTruthTracking(Context& ctx) {
   ACTS_PYTHON_DECLARE_ALGORITHM(
       ActsExamples::TruthVertexFinder, mex, "TruthVertexFinder", inputParticles,
       outputProtoVertices, excludeSecondaries, separateSecondaries);
+
+  ACTS_PYTHON_DECLARE_ALGORITHM(
+      ActsExamples::TrackModifier, mex, "TrackModifier", inputTrajectories,
+      inputTrackParameters, outputTrajectories, outputTrackParameters,
+      dropCovariance, covScale, killTime);
 }
 
 }  // namespace Acts::Python