feat: Seed CSV Writing (#2690)

This PR add a CSV writing mechanism for seed, those can then be used later to train ML based seed filter (once implemented). It also add the seed id to the stored information in the CSV track writing and a python based script that allowed us to matched the best reconstructed track with the corresponding seed.

Examples/Algorithms/TrackFinding/src/TrackFindingAlgorithm.cpp

@@ -157,7 +157,9 @@ ActsExamples::ProcessCode ActsExamples::TrackFindingAlgorithm::execute(
 
     auto& tracksForSeed = result.value();
     for (auto& track : tracksForSeed) {
-      seedNumber(track) = nSeed;
+      // Set the seed number, this number decrease by 1 since the seed number
+      // has already been updated
+      seedNumber(track) = nSeed - 1;
       if (!m_trackSelector.has_value() ||
           m_trackSelector->isValidTrack(track)) {
         auto destProxy = tracks.getTrack(tracks.addTrack());

Examples/Algorithms/TrackFindingExaTrkX/src/TrackFindingFromPrototrackAlgorithm.cpp

@@ -162,7 +162,9 @@ ActsExamples::ProcessCode TrackFindingFromPrototrackAlgorithm::execute(
     nTracksPerSeeds.push_back(tracksForSeed.size());
 
     for (auto& track : tracksForSeed) {
-      seedNumber(track) = nSeed;
+      // Set the seed number, this number decrease by 1 since the seed number
+      // has already been updated
+      seedNumber(track) = nSeed - 1;
     }
   }
 

Examples/Io/Csv/CMakeLists.txt

@@ -14,6 +14,7 @@ add_library(
   src/CsvTrackingGeometryWriter.cpp
   src/CsvTrackParameterReader.cpp
   src/CsvTrackParameterWriter.cpp
+  src/CsvSeedWriter.cpp 
   src/CsvTrackWriter.cpp 
   src/CsvProtoTrackWriter.cpp
   src/CsvSpacePointWriter.cpp

Examples/Io/Csv/include/ActsExamples/Io/Csv/CsvSeedWriter.hpp

@@ -0,0 +1,107 @@
+// 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 "Acts/EventData/MultiTrajectoryHelpers.hpp"
+#include "ActsExamples/EventData/ProtoTrack.hpp"
+#include "ActsExamples/EventData/SimHit.hpp"
+#include "ActsExamples/EventData/SimParticle.hpp"
+#include "ActsExamples/EventData/SimSeed.hpp"
+#include "ActsExamples/EventData/Trajectories.hpp"
+#include "ActsExamples/Framework/DataHandle.hpp"
+#include "ActsExamples/Framework/WriterT.hpp"
+#include "ActsFatras/EventData/Barcode.hpp"
+
+#include <fstream>
+
+using namespace Acts::UnitLiterals;
+
+namespace ActsExamples {
+
+/// @class CsvSeedWriter
+///
+/// Write out the seed reconstructed by the seeding algorithm in
+/// comma-separated-value format.
+///
+/// This writes one file per event into the configured output directory. By
+/// default it writes to the current working directory.
+/// Files are named using the following schema
+///
+///     event000000001-seed.csv
+///     event000000002-seed.csv
+///
+/// and each line in the file corresponds to one seed.
+class CsvSeedWriter : public WriterT<TrackParametersContainer> {
+ public:
+  using HitParticlesMap = IndexMultimap<ActsFatras::Barcode>;
+  using HitSimHitsMap = IndexMultimap<Index>;
+
+  struct Config {
+    /// Input estimated track parameters collection.
+    std::string inputTrackParameters;
+    /// Input seed collection.
+    std::string inputSimSeeds;
+    /// Input collection of simulated hits.
+    std::string inputSimHits;
+    /// Input hit-particles map collection.
+    std::string inputMeasurementParticlesMap;
+    /// Input collection to map measured hits to simulated hits.
+    std::string inputMeasurementSimHitsMap;
+    /// output filename.
+    std::string fileName = "Seed.csv";
+    /// output directory
+    std::string outputDir;
+  };
+
+  /// Constructor
+  ///
+  /// @param config Configuration struct
+  /// @param level Message level declaration
+  CsvSeedWriter(const Config& config,
+                Acts::Logging::Level level = Acts::Logging::INFO);
+
+  /// Get readonly access to the config parameters
+  const Config& config() const { return m_cfg; }
+
+ protected:
+  /// @brief Write method called by the base class
+  /// @param [in] ctx is the algorithm context for event information
+  /// @param [in] trackParams are parameters to write
+  ProcessCode writeT(const AlgorithmContext& ctx,
+                     const TrackParametersContainer& trackParams) override;
+
+ private:
+  Config m_cfg;  ///< The config class
+
+  ReadDataHandle<SimParticleContainer> m_inputParticles{this, "InputParticles"};
+  ReadDataHandle<SimSeedContainer> m_inputSimSeeds{this, "InputSimSeeds"};
+  ReadDataHandle<SimHitContainer> m_inputSimHits{this, "InputSimHits"};
+  ReadDataHandle<HitParticlesMap> m_inputMeasurementParticlesMap{
+      this, "InputMeasurementParticlesMap"};
+  ReadDataHandle<HitSimHitsMap> m_inputMeasurementSimHitsMap{
+      this, "InputMeasurementSimHitsMap"};
+
+  /// @brief Struct for brief seed summary info
+  ///
+  struct SeedInfo {
+    std::size_t seedId = 0;
+    ActsFatras::Barcode particleId;
+    float seedPt = -1;
+    float seedPhi = 0;
+    float seedEta = 0;
+    float vertexZ = 0;
+    float quality = -1;
+    boost::container::small_vector<Acts::Vector3, 3> globalPosition;
+    float truthDistance = -1;
+    std::string seedType = "unknown";
+    ProtoTrack measurementsID;
+  };  // trackInfo struct
+};
+
+}  // namespace ActsExamples

Examples/Io/Csv/include/ActsExamples/Io/Csv/CsvTrackWriter.hpp

@@ -89,6 +89,7 @@ class CsvTrackWriter : public WriterT<ConstTrackContainer> {
   ///
   struct TrackInfo : public Acts::MultiTrajectoryHelpers::TrajectoryState {
     std::size_t trackId = 0;
+    unsigned int seedId = 0;
     ActsFatras::Barcode particleId;
     std::size_t nMajorityHits = 0;
     std::string trackType;

Examples/Io/Csv/src/CsvSeedWriter.cpp

@@ -0,0 +1,202 @@
+// 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/Io/Csv/CsvSeedWriter.hpp"
+
+#include "Acts/EventData/TrackParameters.hpp"
+#include "Acts/Seeding/Seed.hpp"
+#include "Acts/Utilities/Helpers.hpp"
+#include "ActsExamples/EventData/AverageSimHits.hpp"
+#include "ActsExamples/EventData/Index.hpp"
+#include "ActsExamples/EventData/Measurement.hpp"
+#include "ActsExamples/EventData/SimHit.hpp"
+#include "ActsExamples/EventData/SimParticle.hpp"
+#include "ActsExamples/EventData/SimSeed.hpp"
+#include "ActsExamples/Utilities/EventDataTransforms.hpp"
+#include "ActsExamples/Utilities/Paths.hpp"
+#include "ActsExamples/Utilities/Range.hpp"
+#include "ActsExamples/Validation/TrackClassification.hpp"
+
+#include <ios>
+#include <iostream>
+#include <stdexcept>
+#include <string>
+#include <unordered_map>
+#include <unordered_set>
+
+using Acts::VectorHelpers::eta;
+using Acts::VectorHelpers::phi;
+using Acts::VectorHelpers::theta;
+
+ActsExamples::CsvSeedWriter::CsvSeedWriter(
+    const ActsExamples::CsvSeedWriter::Config& config,
+    Acts::Logging::Level level)
+    : WriterT<TrackParametersContainer>(config.inputTrackParameters,
+                                        "CsvSeedWriter", level),
+      m_cfg(config) {
+  if (m_cfg.inputSimSeeds.empty()) {
+    throw std::invalid_argument("Missing space points input collection");
+  }
+  if (m_cfg.inputSimHits.empty()) {
+    throw std::invalid_argument("Missing simulated hits input collection");
+  }
+  if (m_cfg.inputMeasurementParticlesMap.empty()) {
+    throw std::invalid_argument("Missing hit-particles map input collection");
+  }
+  if (m_cfg.inputMeasurementSimHitsMap.empty()) {
+    throw std::invalid_argument(
+        "Missing hit-simulated-hits map input collection");
+  }
+  if (m_cfg.fileName.empty()) {
+    throw std::invalid_argument("Missing output filename");
+  }
+  if (m_cfg.outputDir.empty()) {
+    throw std::invalid_argument("Missing output directory");
+  }
+
+  m_inputSimSeeds.initialize(m_cfg.inputSimSeeds);
+  m_inputSimHits.initialize(m_cfg.inputSimHits);
+  m_inputMeasurementParticlesMap.initialize(m_cfg.inputMeasurementParticlesMap);
+  m_inputMeasurementSimHitsMap.initialize(m_cfg.inputMeasurementSimHitsMap);
+}
+
+ActsExamples::ProcessCode ActsExamples::CsvSeedWriter::writeT(
+    const ActsExamples::AlgorithmContext& ctx,
+    const TrackParametersContainer& trackParams) {
+  // Read additional input collections
+  const auto& seeds = m_inputSimSeeds(ctx);
+  const auto& simHits = m_inputSimHits(ctx);
+  const auto& hitParticlesMap = m_inputMeasurementParticlesMap(ctx);
+  const auto& hitSimHitsMap = m_inputMeasurementSimHitsMap(ctx);
+
+  std::string path =
+      perEventFilepath(m_cfg.outputDir, m_cfg.fileName, ctx.eventNumber);
+
+  std::ofstream mos(path, std::ofstream::out | std::ofstream::trunc);
+  if (!mos) {
+    throw std::ios_base::failure("Could not open '" + path + "' to write");
+  }
+
+  std::unordered_map<std::size_t, SeedInfo> infoMap;
+  std::unordered_map<ActsFatras::Barcode, std::pair<std::size_t, float>>
+      goodSeed;
+
+  // Loop over the estimated track parameters
+  for (std::size_t iparams = 0; iparams < trackParams.size(); ++iparams) {
+    // The estimated bound parameters vector
+    const auto params = trackParams[iparams].parameters();
+
+    float seedPhi = params[Acts::eBoundPhi];
+    float seedEta = std::atanh(std::cos(params[Acts::eBoundTheta]));
+
+    // Get the proto track from which the track parameters are estimated
+    const auto& seed = seeds[iparams];
+    const auto& ptrack = seedToPrototrack(seed);
+
+    std::vector<ParticleHitCount> particleHitCounts;
+    identifyContributingParticles(hitParticlesMap, ptrack, particleHitCounts);
+    bool truthMatched = false;
+    float truthDistance = -1;
+    auto majorityParticleId = particleHitCounts.front().particleId;
+    // Seed are considered truth matched if they have only one contributing
+    // particle
+    if (particleHitCounts.size() == 1) {
+      truthMatched = true;
+      // Get the index of the first space point
+      const auto& hitIdx = ptrack.front();
+      // Get the sim hits via the measurement to sim hits map
+      auto indices = makeRange(hitSimHitsMap.equal_range(hitIdx));
+      // Get the truth particle direction from the sim hits
+      Acts::Vector3 truthUnitDir = {0, 0, 0};
+      for (auto [_, simHitIdx] : indices) {
+        const auto& simHit = *simHits.nth(simHitIdx);
+        if (simHit.particleId() == majorityParticleId) {
+          truthUnitDir = simHit.direction();
+        }
+      }
+      // Compute the distance between the truth and estimated directions
+      float truthPhi = phi(truthUnitDir);
+      float truthEta = std::atanh(std::cos(theta(truthUnitDir)));
+      float dEta = fabs(truthEta - seedEta);
+      float dPhi = fabs(truthPhi - seedPhi) < M_PI
+                       ? fabs(truthPhi - seedPhi)
+                       : fabs(truthPhi - seedPhi) - M_PI;
+      truthDistance = sqrt(dPhi * dPhi + dEta * dEta);
+      // If the seed is truth matched, check if it is the closest one for the
+      // contributing particle
+      if (goodSeed.find(majorityParticleId) != goodSeed.end()) {
+        if (goodSeed[majorityParticleId].second > truthDistance) {
+          goodSeed[majorityParticleId] = std::make_pair(iparams, truthDistance);
+        }
+      } else {
+        goodSeed[majorityParticleId] = std::make_pair(iparams, truthDistance);
+      }
+    }
+    // Store the global position of the space points
+    boost::container::small_vector<Acts::Vector3, 3> globalPosition;
+    for (auto spacePointPtr : seed.sp()) {
+      Acts::Vector3 pos(spacePointPtr->x(), spacePointPtr->y(),
+                        spacePointPtr->z());
+      globalPosition.push_back(pos);
+    }
+
+    // track info
+    SeedInfo toAdd;
+    toAdd.seedId = iparams;
+    toAdd.particleId = majorityParticleId;
+    toAdd.seedPt = std::abs(1.0 / params[Acts::eBoundQOverP]) *
+                   std::sin(params[Acts::eBoundTheta]);
+    toAdd.seedPhi = seedPhi;
+    toAdd.seedEta = seedEta;
+    toAdd.vertexZ = seed.z();
+    toAdd.quality = seed.seedQuality();
+    toAdd.globalPosition = globalPosition;
+    toAdd.truthDistance = truthDistance;
+    toAdd.seedType = truthMatched ? "duplicate" : "fake";
+    toAdd.measurementsID = ptrack;
+
+    infoMap[toAdd.seedId] = toAdd;
+  }
+
+  mos << "seed_id,particleId,"
+      << "pT,eta,phi,"
+      << "bX,bY,bZ,"
+      << "mX,mY,mZ,"
+      << "tX,tY,tZ,"
+      << "good/duplicate/fake,"
+      << "vertexZ,quality,"
+      << "Hits_ID" << '\n';
+
+  for (auto& [id, info] : infoMap) {
+    if (goodSeed[info.particleId].first == id) {
+      info.seedType = "good";
+    }
+    // write the track info
+    mos << info.seedId << ",";
+    mos << info.particleId << ",";
+    mos << info.seedPt << ",";
+    mos << info.seedEta << ",";
+    mos << info.seedPhi << ",";
+    for (auto& point : info.globalPosition) {
+      mos << point.x() << ",";
+      mos << point.y() << ",";
+      mos << point.z() << ",";
+    }
+    mos << info.seedType << ",";
+    mos << info.vertexZ << ",";
+    mos << info.quality << ",";
+    mos << "\"[";
+    for (auto& ID : info.measurementsID) {
+      mos << ID << ",";
+    }
+    mos << "]\"";
+    mos << '\n';
+  }
+
+  return ProcessCode::SUCCESS;
+}

Examples/Io/Csv/src/CsvTrackWriter.cpp

@@ -108,9 +108,18 @@ ProcessCode CsvTrackWriter::writeT(const AlgorithmContext& context,
       // n Majority hits
       nMajorityHits = particleHitCount.front().hitCount;
     }
+
+    static const Acts::ConstTrackAccessor<unsigned int> seedNumber(
+        "trackGroup");
+
     // track info
     TrackInfo toAdd;
     toAdd.trackId = trackId;
+    if (tracks.hasColumn(Acts::hashString("trackGroup"))) {
+      toAdd.seedId = seedNumber(track);
+    } else {
+      toAdd.seedId = 0;
+    }
     toAdd.particleId = majorityParticleId;
     toAdd.nStates = track.nTrackStates();
     toAdd.nMajorityHits = nMajorityHits;
@@ -166,7 +175,7 @@ ProcessCode CsvTrackWriter::writeT(const AlgorithmContext& context,
   }
 
   // write csv header
-  mos << "track_id,particleId,"
+  mos << "track_id,seed_id,particleId,"
       << "nStates,nMajorityHits,nMeasurements,nOutliers,nHoles,nSharedHits,"
       << "chi2,ndf,chi2/ndf,"
       << "pT,eta,phi,"
@@ -190,6 +199,7 @@ ProcessCode CsvTrackWriter::writeT(const AlgorithmContext& context,
 
     // write the track info
     mos << trajState.trackId << ",";
+    mos << trajState.seedId << ",";
     mos << trajState.particleId << ",";
     mos << trajState.nStates << ",";
     mos << trajState.nMajorityHits << ",";