feat: Read CSV clusters and measurement-particles-map & refactor

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

@@ -15,6 +15,8 @@
 #include "ActsExamples/EventData/Index.hpp"
 #include "ActsExamples/EventData/IndexSourceLink.hpp"
 #include "ActsExamples/EventData/Measurement.hpp"
+#include "ActsExamples/EventData/SimHit.hpp"
+#include "ActsExamples/EventData/SimParticle.hpp"
 #include "ActsExamples/Framework/DataHandle.hpp"
 #include "ActsExamples/Framework/IReader.hpp"
 #include "ActsExamples/Framework/ProcessCode.hpp"
@@ -62,6 +64,12 @@ class CsvMeasurementReader final : public IReader {
     std::string outputSourceLinks;
     /// Output cluster collection (optional).
     std::string outputClusters;
+
+    /// Input SimHits for measurment-particle map (optional)
+    std::string inputSimHits;
+    /// Output  measurement to particle collection (optional)
+    /// @note Only filled if inputSimHits is given
+    std::string outputMeasurementParticlesMap;
   };
 
   /// Construct the cluster reader.
@@ -98,6 +106,11 @@ class CsvMeasurementReader final : public IReader {
       this, "OutputSourceLinks"};
 
   WriteDataHandle<ClusterContainer> m_outputClusters{this, "OutputClusters"};
+
+  WriteDataHandle<IndexMultimap<ActsFatras::Barcode>>
+      m_outputMeasurementParticlesMap{this, "OutputMeasurementParticlesMap"};
+
+  ReadDataHandle<SimHitContainer> m_inputHits{this, "InputHits"};
 };
 
 }  // namespace ActsExamples

Examples/Io/Csv/src/CsvMeasurementReader.cpp

@@ -37,7 +37,6 @@ ActsExamples::CsvMeasurementReader::CsvMeasurementReader(
     const ActsExamples::CsvMeasurementReader::Config& config,
     Acts::Logging::Level level)
     : m_cfg(config),
-      // TODO check that all files (hits,cells,truth) exists
       m_eventsRange(
           determineEventFilesRange(m_cfg.inputDir, "measurements.csv")),
       m_logger(Acts::getDefaultLogger("CsvMeasurementReader", level)) {
@@ -49,6 +48,24 @@ ActsExamples::CsvMeasurementReader::CsvMeasurementReader(
   m_outputMeasurementSimHitsMap.initialize(m_cfg.outputMeasurementSimHitsMap);
   m_outputSourceLinks.initialize(m_cfg.outputSourceLinks);
   m_outputClusters.maybeInitialize(m_cfg.outputClusters);
+  m_outputMeasurementParticlesMap.maybeInitialize(
+      m_cfg.outputMeasurementParticlesMap);
+  m_inputHits.maybeInitialize(m_cfg.inputSimHits);
+
+  // Check if event ranges match (should also catch missing files)
+  auto checkRange = [&](const std::string& fileStem) {
+    const auto hitmapRange = determineEventFilesRange(m_cfg.inputDir, fileStem);
+    if (hitmapRange.first > m_eventsRange.first or
+        hitmapRange.second < m_eventsRange.second) {
+      throw std::runtime_error("event range mismatch for 'event**-" + fileStem +
+                               "'");
+    }
+  };
+
+  checkRange("measurement-simhit-map.csv");
+  if (not m_cfg.outputClusters.empty()) {
+    checkRange("cells.csv");
+  }
 }
 
 std::string ActsExamples::CsvMeasurementReader::CsvMeasurementReader::name()
@@ -112,14 +129,54 @@ std::vector<ActsExamples::MeasurementData> readMeasurementsByGeometryId(
   return measurements;
 }
 
-std::vector<ActsExamples::CellData> readCellsByHitId(
-    const std::string& inputDir, size_t event) {
-  // timestamp is an optional element
-  auto cells = readEverything<ActsExamples::CellData>(inputDir, "cells.csv",
-                                                      {"timestamp"}, event);
-  // sort for fast hit id look up
-  std::sort(cells.begin(), cells.end(), CompareHitId{});
-  return cells;
+ActsExamples::ClusterContainer makeClusters(
+    const std::unordered_multimap<std::size_t, ActsExamples::CellData>&
+        cellDataMap,
+    std::size_t nMeasurments) {
+  using namespace ActsExamples;
+  ClusterContainer clusters;
+
+  for (auto index = 0ul; index < nMeasurments; ++index) {
+    auto [begin, end] = cellDataMap.equal_range(index);
+
+    // Fill the channels with the iterators
+    Cluster cluster;
+    cluster.channels.reserve(std::distance(begin, end));
+
+    for (auto it = begin; it != end; ++it) {
+      const auto& cellData = it->second;
+      ActsFatras::Channelizer::Segment2D dummySegment = {Acts::Vector2::Zero(),
+                                                         Acts::Vector2::Zero()};
+
+      ActsFatras::Channelizer::Bin2D bin{
+          static_cast<unsigned int>(cellData.channel0),
+          static_cast<unsigned int>(cellData.channel1)};
+
+      cluster.channels.emplace_back(bin, dummySegment, cellData.value);
+    }
+
+    // update the iterator
+
+    // Compute cluster size
+    if (not cluster.channels.empty()) {
+      auto compareX = [](const auto& a, const auto& b) {
+        return a.bin[0] < b.bin[0];
+      };
+      auto compareY = [](const auto& a, const auto& b) {
+        return a.bin[1] < b.bin[1];
+      };
+
+      auto [minX, maxX] = std::minmax_element(cluster.channels.begin(),
+                                              cluster.channels.end(), compareX);
+      auto [minY, maxY] = std::minmax_element(cluster.channels.begin(),
+                                              cluster.channels.end(), compareY);
+      cluster.sizeLoc0 = 1 + maxX->bin[0] - minX->bin[0];
+      cluster.sizeLoc1 = 1 + maxY->bin[1] - minY->bin[1];
+    }
+
+    clusters.push_back(cluster);
+  }
+  return clusters;
 }
 
 }  // namespace
@@ -136,14 +193,8 @@ ActsExamples::ProcessCode ActsExamples::CsvMeasurementReader::read(
   auto measurementData =
       readMeasurementsByGeometryId(m_cfg.inputDir, ctx.eventNumber);
 
-  std::vector<ActsExamples::CellData> cellData = {};
-  if (not m_cfg.outputClusters.empty()) {
-    cellData = readCellsByHitId(m_cfg.inputDir, ctx.eventNumber);
-  }
-
   // Prepare containers for the hit data using the framework event data types
   GeometryIdMultimap<Measurement> orderedMeasurements;
-  ClusterContainer clusters;
   IndexMultimap<Index> measurementSimHitsMap;
   IndexSourceLinkContainer sourceLinks;
   // need list here for stable addresses
@@ -199,8 +250,8 @@ ActsExamples::ProcessCode ActsExamples::CsvMeasurementReader::read(
 
     // The measurement container is unordered and the index under which
     // the measurement will be stored is known before adding it.
-    Index hitIdx = orderedMeasurements.size();
-    IndexSourceLink& sourceLink = sourceLinkStorage.emplace_back(geoId, hitIdx);
+    const Index index = orderedMeasurements.size();
+    IndexSourceLink& sourceLink = sourceLinkStorage.emplace_back(geoId, index);
     auto measurement = createMeasurement(dParameters, sourceLink);
 
     // Due to the previous sorting of the raw hit data by geometry id, new
@@ -222,13 +273,68 @@ ActsExamples::ProcessCode ActsExamples::CsvMeasurementReader::read(
     measurements.emplace_back(std::move(meas));
   }
 
+  // Generate measurment-particles-map
+  if (m_inputHits.isInitialized() &&
+      m_outputMeasurementParticlesMap.isInitialized()) {
+    const auto hits = m_inputHits(ctx);
+
+    IndexMultimap<ActsFatras::Barcode> outputMap;
+
+    for (const auto& [measIdx, hitIdx] : measurementSimHitsMap) {
+      const auto& hit = hits.nth(hitIdx);
+      outputMap.emplace(measIdx, hit->particleId());
+    }
+
+    m_outputMeasurementParticlesMap(ctx, std::move(outputMap));
+  }
+
   // Write the data to the EventStore
   m_outputMeasurements(ctx, std::move(measurements));
   m_outputMeasurementSimHitsMap(ctx, std::move(measurementSimHitsMap));
   m_outputSourceLinks(ctx, std::move(sourceLinks));
-  if (not clusters.empty()) {
-    m_outputClusters(ctx, std::move(clusters));
+
+  /////////////////////////
+  // Cluster information //
+  /////////////////////////
+
+  if (m_cfg.outputClusters.empty()) {
+    return ActsExamples::ProcessCode::SUCCESS;
   }
 
+  std::vector<ActsExamples::CellData> cellData;
+
+  // This allows seamless import of files created with a older version where
+  // the measurment_id-column is still named hit_id
+  try {
+    cellData = readEverything<ActsExamples::CellData>(
+        m_cfg.inputDir, "cells.csv", {"timestamp"}, ctx.eventNumber);
+  } catch (std::runtime_error& e) {
+    // Rethrow exception if it is not about the measurement_id-column
+    if (std::string(e.what()).find("Missing header column 'measurement_id'") ==
+        std::string::npos) {
+      throw;
+    }
+
+    const auto oldCellData = readEverything<ActsExamples::CellDataLegacy>(
+        m_cfg.inputDir, "cells.csv", {"timestamp"}, ctx.eventNumber);
+
+    auto fromLegacy = [](const CellDataLegacy& old) {
+      return CellData{old.geometry_id, old.hit_id,    old.channel0,
+                      old.channel1,    old.timestamp, old.value};
+    };
+
+    cellData.resize(oldCellData.size());
+    std::transform(oldCellData.begin(), oldCellData.end(), cellData.begin(),
+                   fromLegacy);
+  }
+
+  std::unordered_multimap<std::size_t, ActsExamples::CellData> cellDataMap;
+  for (const auto& cd : cellData) {
+    cellDataMap.emplace(cd.measurement_id, cd);
+  }
+
+  auto clusters = makeClusters(cellDataMap, orderedMeasurements.size());
+  m_outputClusters(ctx, std::move(clusters));
+
   return ActsExamples::ProcessCode::SUCCESS;
 }

Examples/Io/Csv/src/CsvMeasurementWriter.cpp

@@ -83,18 +83,18 @@ ActsExamples::ProcessCode ActsExamples::CsvMeasurementWriter::writeT(
   MeasurementData meas;
   CellData cell;
 
-  // Will be reused as hit counter
+  // Will be reused as measurement counter
   meas.measurement_id = 0;
 
   ACTS_VERBOSE("Writing " << measurements.size()
                           << " measurements in this event.");
 
-  for (Index hitIdx = 0u; hitIdx < measurements.size(); ++hitIdx) {
-    const auto& measurement = measurements[hitIdx];
+  for (Index measIdx = 0u; measIdx < measurements.size(); ++measIdx) {
+    const auto& measurement = measurements[measIdx];
 
-    auto simHitIndices = makeRange(measurementSimHitsMap.equal_range(hitIdx));
+    auto simHitIndices = makeRange(measurementSimHitsMap.equal_range(measIdx));
     for (auto [_, simHitIdx] : simHitIndices) {
-      writerMeasurementSimHitMap.append({hitIdx, simHitIdx});
+      writerMeasurementSimHitMap.append({measIdx, simHitIdx});
     }
 
     std::visit(
@@ -131,9 +131,9 @@ ActsExamples::ProcessCode ActsExamples::CsvMeasurementWriter::writeT(
 
           // CLUSTER / channel information ------------------------------
           if (not clusters.empty() && writerCells) {
-            auto cluster = clusters[hitIdx];
+            auto cluster = clusters[measIdx];
             cell.geometry_id = meas.geometry_id;
-            cell.hit_id = meas.measurement_id;
+            cell.measurement_id = meas.measurement_id;
             for (auto& c : cluster.channels) {
               cell.channel0 = c.bin[0];
               cell.channel1 = c.bin[1];

Examples/Io/Csv/src/CsvOutputData.hpp

@@ -139,9 +139,28 @@ struct MeasurementData {
 struct CellData {
   /// Hit surface identifier.
   uint64_t geometry_id = 0u;
-  /// Event-unique hit identifier. As defined for the simulated hit above and
-  /// used to link back to it; same value can appear multiple times for clusters
-  /// with more than one active cell.
+  /// Event-unique measurement identifier. As defined for the measurement above
+  /// and used to link back to it; same value can appear multiple times for
+  /// clusters with more than one active cell.
+  uint64_t measurement_id = 0;
+  /// Digital cell address/ channel
+  int32_t channel0 = 0, channel1 = 0;
+  /// Digital cell timestamp. Not available in the TrackML datasets.
+  float timestamp = 0;
+  /// (Digital) measured cell value, e.g. amplitude or time-over-threshold.
+  float value = 0;
+
+  DFE_NAMEDTUPLE(CellData, geometry_id, measurement_id, channel0, channel1,
+                 timestamp, value);
+};
+
+// uses hit id
+struct CellDataLegacy {
+  /// Hit surface identifier.
+  uint64_t geometry_id = 0u;
+  /// Event-unique measurement identifier. As defined for the measurement above
+  /// and used to link back to it; same value can appear multiple times for
+  /// clusters with more than one active cell.
   uint64_t hit_id = 0;
   /// Digital cell address/ channel
   int32_t channel0 = 0, channel1 = 0;
@@ -150,8 +169,8 @@ struct CellData {
   /// (Digital) measured cell value, e.g. amplitude or time-over-threshold.
   float value = 0;
 
-  DFE_NAMEDTUPLE(CellData, geometry_id, hit_id, channel0, channel1, timestamp,
-                 value);
+  DFE_NAMEDTUPLE(CellDataLegacy, geometry_id, hit_id, channel0, channel1,
+                 timestamp, value);
 };
 
 struct SurfaceData {

Examples/Io/Csv/src/CsvPlanarClusterReader.cpp

@@ -137,11 +137,11 @@ std::vector<ActsExamples::HitData> readHitsByGeometryId(
   return hits;
 }
 
-std::vector<ActsExamples::CellData> readCellsByHitId(
+std::vector<ActsExamples::CellDataLegacy> readCellsByHitId(
     const std::string& inputDir, size_t event) {
   // timestamp is an optional element
-  auto cells = readEverything<ActsExamples::CellData>(inputDir, "cells.csv",
-                                                      {"timestamp"}, event);
+  auto cells = readEverything<ActsExamples::CellDataLegacy>(
+      inputDir, "cells.csv", {"timestamp"}, event);
   // sort for fast hit id look up
   std::sort(cells.begin(), cells.end(), CompareHitId{});
   return cells;

Examples/Io/Csv/src/CsvPlanarClusterWriter.cpp

@@ -65,13 +65,13 @@ ActsExamples::ProcessCode ActsExamples::CsvPlanarClusterWriter::writeT(
       perEventFilepath(m_cfg.outputDir, "truth.csv", ctx.eventNumber);
 
   dfe::NamedTupleCsvWriter<HitData> writerHits(pathHits, m_cfg.outputPrecision);
-  dfe::NamedTupleCsvWriter<CellData> writerCells(pathCells,
-                                                 m_cfg.outputPrecision);
+  dfe::NamedTupleCsvWriter<CellDataLegacy> writerCells(pathCells,
+                                                       m_cfg.outputPrecision);
   dfe::NamedTupleCsvWriter<TruthHitData> writerTruth(pathTruth,
                                                      m_cfg.outputPrecision);
 
   HitData hit;
-  CellData cell;
+  CellDataLegacy cell;
   TruthHitData truth;
   // will be reused as hit counter
   hit.hit_id = 0;

Examples/Python/src/Input.cpp

@@ -55,10 +55,10 @@ void addInput(Context& ctx) {
                              "CsvParticleReader", inputDir, inputStem,
                              outputParticles);
 
-  ACTS_PYTHON_DECLARE_READER(ActsExamples::CsvMeasurementReader, mex,
-                             "CsvMeasurementReader", inputDir,
-                             outputMeasurements, outputMeasurementSimHitsMap,
-                             outputSourceLinks, outputClusters);
+  ACTS_PYTHON_DECLARE_READER(
+      ActsExamples::CsvMeasurementReader, mex, "CsvMeasurementReader", inputDir,
+      outputMeasurements, outputMeasurementSimHitsMap, outputSourceLinks,
+      outputClusters, outputMeasurementParticlesMap, inputSimHits);
 
   ACTS_PYTHON_DECLARE_READER(ActsExamples::CsvPlanarClusterReader, mex,
                              "CsvPlanarClusterReader", inputDir, outputClusters,