feat: RootTrackSummaryWriter can write multiple track collections

Examples/Io/Root/include/ActsExamples/Io/Root/RootTrackSummaryWriter.hpp

@@ -9,13 +9,12 @@
 #pragma once
 
 #include "Acts/Utilities/Logger.hpp"
-#include "ActsExamples/EventData/SimHit.hpp"
 #include "ActsExamples/EventData/SimParticle.hpp"
 #include "ActsExamples/EventData/Track.hpp"
 #include "ActsExamples/EventData/TruthMatching.hpp"
 #include "ActsExamples/Framework/DataHandle.hpp"
+#include "ActsExamples/Framework/IWriter.hpp"
 #include "ActsExamples/Framework/ProcessCode.hpp"
-#include "ActsExamples/Framework/WriterT.hpp"
 
 #include <cstdint>
 #include <mutex>
@@ -44,11 +43,11 @@ namespace ActsExamples {
 /// done by setting the Config::rootFile pointer to an existing file.
 ///
 /// Safe to use from multiple writer threads - uses a std::mutex lock.
-class RootTrackSummaryWriter final : public WriterT<ConstTrackContainer> {
+class RootTrackSummaryWriter final : public IWriter {
  public:
   struct Config {
     /// Input (fitted) tracks collection
-    std::string inputTracks;
+    std::vector<std::string> inputTrackContainers;
     /// Input particles collection (optional).
     std::string inputParticles;
     /// Input track-particle matching (optional).
@@ -80,21 +79,31 @@ class RootTrackSummaryWriter final : public WriterT<ConstTrackContainer> {
   /// Get readonly access to the config parameters
   const Config& config() const { return m_cfg; }
 
+  ProcessCode write(const AlgorithmContext& ctx) override;
+
+  std::string name() const override { return "RootTrackSummaryWriter"; };
+
  protected:
-  /// @brief Write method called by the base class
-  /// @param [in] ctx is the algorithm context for event information
-  /// @param [in] tracks are what to be written out
-  ProcessCode writeT(const AlgorithmContext& ctx,
-                     const ConstTrackContainer& tracks) override;
+  ProcessCode write(const AlgorithmContext& ctx,
+                    const TrackParticleMatching& trackParticleMatching,
+                    const SimParticleContainer& particles,
+                    const ConstTrackContainer& tracks);
 
  private:
   /// The config class
   Config m_cfg;
 
+  std::unique_ptr<const Acts::Logger> m_logger;
+
+  const Acts::Logger& logger() const { return *m_logger; }
+
   ReadDataHandle<SimParticleContainer> m_inputParticles{this, "InputParticles"};
   ReadDataHandle<TrackParticleMatching> m_inputTrackParticleMatching{
       this, "InputTrackParticleMatching"};
 
+  std::vector<std::unique_ptr<ReadDataHandle<ConstTrackContainer>>>
+      m_inputTrackContainers{};
+
   /// Mutex used to protect multi-threaded writes
   std::mutex m_writeMutex;
   /// The output file
@@ -172,6 +181,7 @@ class RootTrackSummaryWriter final : public WriterT<ConstTrackContainer> {
   std::vector<float> m_t_z0;
   /// Production radius of majority particle
   std::vector<float> m_t_prodR;
+  std::vector<unsigned int> m_t_pdg;
 
   /// If the track has fitted parameter
   std::vector<bool> m_hasFittedParams;
@@ -228,6 +238,8 @@ class RootTrackSummaryWriter final : public WriterT<ConstTrackContainer> {
   /// Fitted parameters eT pull of track
   std::vector<float> m_pull_eT_fit;
 
+  std::vector<unsigned int> m_hypo_pdg;
+
   // entries of the full covariance matrix. One block for every row of the
   // matrix
   std::vector<float> m_cov_eLOC0_eLOC0;

Examples/Io/Root/src/RootTrackSummaryWriter.cpp

@@ -15,12 +15,15 @@
 #include "Acts/Surfaces/Surface.hpp"
 #include "Acts/TrackFitting/GsfOptions.hpp"
 #include "Acts/Utilities/Intersection.hpp"
+#include "Acts/Utilities/Logger.hpp"
 #include "Acts/Utilities/MultiIndex.hpp"
 #include "Acts/Utilities/Result.hpp"
 #include "Acts/Utilities/detail/periodic.hpp"
 #include "ActsExamples/EventData/IndexSourceLink.hpp"
+#include "ActsExamples/EventData/Track.hpp"
 #include "ActsExamples/EventData/TruthMatching.hpp"
 #include "ActsExamples/Framework/AlgorithmContext.hpp"
+#include "ActsExamples/Framework/ProcessCode.hpp"
 #include "ActsExamples/Framework/WriterT.hpp"
 #include "ActsExamples/Validation/TrackClassification.hpp"
 #include "ActsFatras/EventData/Barcode.hpp"
@@ -50,8 +53,7 @@ namespace ActsExamples {
 
 RootTrackSummaryWriter::RootTrackSummaryWriter(
     const RootTrackSummaryWriter::Config& config, Acts::Logging::Level level)
-    : WriterT(config.inputTracks, "RootTrackSummaryWriter", level),
-      m_cfg(config) {
+    : m_cfg(config), m_logger(Acts::getDefaultLogger(name(), level)) {
   // tracks collection name is already checked by base ctor
   if (m_cfg.filePath.empty()) {
     throw std::invalid_argument("Missing output filename");
@@ -64,6 +66,18 @@ RootTrackSummaryWriter::RootTrackSummaryWriter(
   m_inputTrackParticleMatching.maybeInitialize(
       m_cfg.inputTrackParticleMatching);
 
+  for (const auto& name : m_cfg.inputTrackContainers) {
+    if (name.empty()) {
+      throw std::invalid_argument("Invalid track input collection");
+    }
+
+    auto& handle = m_inputTrackContainers.emplace_back(
+        std::make_unique<ReadDataHandle<ConstTrackContainer>>(
+            this, "InputTracksCollection#" +
+                      std::to_string(m_inputTrackContainers.size())));
+    handle->initialize(name);
+  }
+
   // Setup ROOT I/O
   auto path = m_cfg.filePath;
   m_outputFile = TFile::Open(path.c_str(), m_cfg.fileMode.c_str());
@@ -113,6 +127,7 @@ RootTrackSummaryWriter::RootTrackSummaryWriter(
   m_outputTree->Branch("t_d0", &m_t_d0);
   m_outputTree->Branch("t_z0", &m_t_z0);
   m_outputTree->Branch("t_prodR", &m_t_prodR);
+  m_outputTree->Branch("t_pdg", &m_t_pdg);
 
   m_outputTree->Branch("hasFittedParams", &m_hasFittedParams);
   m_outputTree->Branch("eLOC0_fit", &m_eLOC0_fit);
@@ -139,6 +154,7 @@ RootTrackSummaryWriter::RootTrackSummaryWriter(
   m_outputTree->Branch("pull_eTHETA_fit", &m_pull_eTHETA_fit);
   m_outputTree->Branch("pull_eQOP_fit", &m_pull_eQOP_fit);
   m_outputTree->Branch("pull_eT_fit", &m_pull_eT_fit);
+  m_outputTree->Branch("hypo_pdg", &m_hypo_pdg);
 
   if (m_cfg.writeGsfSpecific) {
     m_outputTree->Branch("max_material_fwd", &m_gsf_max_material_fwd);
@@ -214,28 +230,10 @@ ProcessCode RootTrackSummaryWriter::finalize() {
   return ProcessCode::SUCCESS;
 }
 
-ProcessCode RootTrackSummaryWriter::writeT(const AlgorithmContext& ctx,
-                                           const ConstTrackContainer& tracks) {
-  // In case we do not have truth info, we bind to a empty collection
-  const static SimParticleContainer emptyParticles;
-  const static TrackParticleMatching emptyTrackParticleMatching;
-
-  const auto& particles =
-      m_inputParticles.isInitialized() ? m_inputParticles(ctx) : emptyParticles;
-  const auto& trackParticleMatching =
-      m_inputTrackParticleMatching.isInitialized()
-          ? m_inputTrackParticleMatching(ctx)
-          : emptyTrackParticleMatching;
-
-  // For each particle within a track, how many hits did it contribute
-  std::vector<ParticleHitCount> particleHitCounts;
-
-  // Exclusive access to the tree while writing
-  std::lock_guard<std::mutex> lock(m_writeMutex);
-
-  // Get the event number
-  m_eventNr = ctx.eventNumber;
-
+ProcessCode RootTrackSummaryWriter::write(
+    const AlgorithmContext& ctx,
+    const TrackParticleMatching& trackParticleMatching,
+    const SimParticleContainer& particles, const ConstTrackContainer& tracks) {
   for (const auto& track : tracks) {
     m_trackNr.push_back(track.index());
 
@@ -301,6 +299,7 @@ ProcessCode RootTrackSummaryWriter::writeT(const AlgorithmContext& ctx,
     float t_z0 = NaNfloat;
     float t_qop = NaNfloat;
     float t_prodR = NaNfloat;
+    unsigned int t_pdg = 0;
 
     // Get the perigee surface
     const Acts::Surface* pSurface =
@@ -342,6 +341,7 @@ ProcessCode RootTrackSummaryWriter::writeT(const AlgorithmContext& ctx,
         t_pT = t_p * perp(particle.direction());
         t_qop = particle.qOverP();
         t_prodR = std::sqrt(t_vx * t_vx + t_vy * t_vy);
+        t_pdg = particle.pdg();
 
         if (pSurface != nullptr) {
           auto intersection =
@@ -394,6 +394,7 @@ ProcessCode RootTrackSummaryWriter::writeT(const AlgorithmContext& ctx,
     m_t_d0.push_back(t_d0);
     m_t_z0.push_back(t_z0);
     m_t_prodR.push_back(t_prodR);
+    m_t_pdg.push_back(t_pdg);
 
     // Initialize the fitted track parameters info
     std::array<float, Acts::eBoundSize> param = {NaNfloat, NaNfloat, NaNfloat,
@@ -465,6 +466,7 @@ ProcessCode RootTrackSummaryWriter::writeT(const AlgorithmContext& ctx,
     m_pull_eTHETA_fit.push_back(pull[Acts::eBoundTheta]);
     m_pull_eQOP_fit.push_back(pull[Acts::eBoundQOverP]);
     m_pull_eT_fit.push_back(pull[Acts::eBoundTime]);
+    m_hypo_pdg.push_back(track.particleHypothesis().absolutePdg());
 
     m_hasFittedParams.push_back(hasFittedParams);
 
@@ -543,6 +545,38 @@ ProcessCode RootTrackSummaryWriter::writeT(const AlgorithmContext& ctx,
     }
   }
 
+  return ProcessCode::SUCCESS;
+}
+
+ProcessCode RootTrackSummaryWriter::write(const AlgorithmContext& ctx) {
+  // In case we do not have truth info, we bind to a empty collection
+  const static SimParticleContainer emptyParticles;
+  const static TrackParticleMatching emptyTrackParticleMatching;
+
+  const auto& particles =
+      m_inputParticles.isInitialized() ? m_inputParticles(ctx) : emptyParticles;
+  const auto& trackParticleMatching =
+      m_inputTrackParticleMatching.isInitialized()
+          ? m_inputTrackParticleMatching(ctx)
+          : emptyTrackParticleMatching;
+
+  // For each particle within a track, how many hits did it contribute
+  std::vector<ParticleHitCount> particleHitCounts;
+
+  // Exclusive access to the tree while writing
+  std::lock_guard<std::mutex> lock(m_writeMutex);
+
+  // Get the event number
+  m_eventNr = ctx.eventNumber;
+
+  for (const auto& handle : m_inputTrackContainers) {
+    const auto& tracks = (*handle)(ctx);
+    if (auto pc = write(ctx, trackParticleMatching, particles, tracks);
+        pc != ProcessCode::SUCCESS) {
+      return pc;
+    }
+  }
+
   // fill the variables
   m_outputTree->Fill();
 

Examples/Python/include/Acts/Plugins/Python/Utilities.hpp

@@ -121,12 +121,6 @@ void patchKwargsConstructor(T& c) {
                  py::arg("config"), py::arg("level"))                       \
             .def_property_readonly("config", &Writer::config);              \
                                                                             \
-    constexpr bool has_write_method =                                       \
-        Acts::Concepts::has_write_method<Writer>;                           \
-                                                                            \
-    if constexpr (has_write_method) {                                       \
-      w.def("write", &Writer::write);                                       \
-    }                                                                       \
     auto c = py::class_<Config>(w, "Config").def(py::init<>());             \
     ACTS_PYTHON_STRUCT(c, __VA_ARGS__);                                     \
   } while (0)

Examples/Python/src/RootOutput.cpp

@@ -199,10 +199,11 @@ void addRootOutput(Context& ctx) {
       inputTracks, inputParticles, inputTrackParticleMatching, inputSimHits,
       inputMeasurementSimHitsMap, filePath, treeName, fileMode);
 
-  ACTS_PYTHON_DECLARE_WRITER(
-      ActsExamples::RootTrackSummaryWriter, mex, "RootTrackSummaryWriter",
-      inputTracks, inputParticles, inputTrackParticleMatching, filePath,
-      treeName, fileMode, writeCovMat, writeGsfSpecific, writeGx2fSpecific);
+  ACTS_PYTHON_DECLARE_WRITER(ActsExamples::RootTrackSummaryWriter, mex,
+                             "RootTrackSummaryWriter", inputTrackContainers,
+                             inputParticles, inputTrackParticleMatching,
+                             filePath, treeName, fileMode, writeCovMat,
+                             writeGsfSpecific, writeGx2fSpecific);
 
   ACTS_PYTHON_DECLARE_WRITER(
       ActsExamples::VertexNTupleWriter, mex, "VertexNTupleWriter",