fix: Root readers ensure event ordering is consistent (#837)

The current Root readers read the event from the root file assuming that the event `eventId` is the entry `eventId` in the tree. This is OK if only one root file is being read, but in case two different root files are read in, objects from two different events might be processed if the root files are written in multi-threaded modes.

This PR fixes the issue by finding the entry that has the desirable event id first. This might be expensive and could be improved by a better solution in the future.

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

@@ -37,7 +37,8 @@ class RootMaterialTrackReader : public IReader {
     std::string treeName = "material-tracks";  ///< name of the output tree
     std::vector<std::string> fileList;         ///< The name of the input file
 
-    unsigned int batchSize = 1;  ///!< The number of tracks per event
+    /// Whether the events are ordered or not
+    bool orderedEvents = true;
 
     /// The default logger
     std::shared_ptr<const Acts::Logger> logger;
@@ -88,6 +89,13 @@ class RootMaterialTrackReader : public IReader {
   /// The input tree name
   TChain* m_inputChain = nullptr;
 
+  /// Event identifier.
+  uint32_t m_eventId;
+
+  /// The entry numbers for accessing events in increased order (there could be
+  /// multiple entries corresponding to one event number)
+  std::vector<long long> m_entryNumbers = {};
+
   float m_v_x;    ///< start global x
   float m_v_y;    ///< start global y
   float m_v_z;    ///< start global z

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

@@ -91,6 +91,9 @@ class RootMaterialTrackWriter
   /// The output tree name
   TTree* m_outputTree;
 
+  /// Event identifier.
+  uint32_t m_eventId;
+
   float m_v_x;    ///< start global x
   float m_v_y;    ///< start global y
   float m_v_z;    ///< start global z

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

@@ -38,6 +38,9 @@ class RootParticleReader : public IReader {
     std::string inputFile;          ///< The name of the input file
     std::string inputDir;           ///< The name of the input dir
 
+    /// Whether the events are ordered or not
+    bool orderedEvents = true;
+
     /// The default logger
     std::shared_ptr<const Acts::Logger> logger;
 
@@ -90,6 +93,10 @@ class RootParticleReader : public IReader {
   /// Event identifier.
   uint32_t m_eventId;
 
+  /// The entry numbers for accessing events in increased order (there could be
+  /// multiple entries corresponding to one event number)
+  std::vector<long long> m_entryNumbers = {};
+
   std::vector<uint64_t>* m_particleId = new std::vector<uint64_t>;
   std::vector<int32_t>* m_particleType = new std::vector<int32_t>;
   std::vector<uint32_t>* m_process = new std::vector<uint32_t>;

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

@@ -37,6 +37,9 @@ class RootTrajectoryParametersReader : public IReader {
     std::string inputFile;  ///< The name of the input file
     std::string inputDir;   ///< The name of the input dir
 
+    /// Whether the events are ordered or not
+    bool orderedEvents = true;
+
     /// The default logger
     std::shared_ptr<const Acts::Logger> logger;
 
@@ -93,6 +96,10 @@ class RootTrajectoryParametersReader : public IReader {
       new std::vector<unsigned int>;  ///< the multi-trajectory sub-trajectory
                                       ///< number
 
+  /// The entry numbers for accessing events in increased order (there could be
+  /// multiple entries corresponding to one event number)
+  std::vector<long long> m_entryNumbers = {};
+
   std::vector<unsigned long>* m_t_barcode =
       new std::vector<long unsigned int>;  ///< Truth particle barcode
   std::vector<int>* m_t_charge =

Examples/Io/Root/src/RootMaterialTrackReader.cpp

@@ -14,13 +14,15 @@
 
 #include <TChain.h>
 #include <TFile.h>
+#include <TMath.h>
 
 ActsExamples::RootMaterialTrackReader::RootMaterialTrackReader(
     const ActsExamples::RootMaterialTrackReader::Config& cfg)
     : ActsExamples::IReader(), m_cfg(cfg), m_events(0), m_inputChain(nullptr) {
   m_inputChain = new TChain(m_cfg.treeName.c_str());
 
   // Set the branches
+  m_inputChain->SetBranchAddress("event_id", &m_eventId);
   m_inputChain->SetBranchAddress("v_x", &m_v_x);
   m_inputChain->SetBranchAddress("v_y", &m_v_y);
   m_inputChain->SetBranchAddress("v_z", &m_v_z);
@@ -54,6 +56,15 @@ ActsExamples::RootMaterialTrackReader::RootMaterialTrackReader(
 
   m_events = m_inputChain->GetEntries();
   ACTS_DEBUG("The full chain has " << m_events << " entries.");
+
+  // If the events are not in order, get the entry numbers for ordered events
+  if (not m_cfg.orderedEvents) {
+    m_entryNumbers.resize(m_events);
+    m_inputChain->Draw("event_id", "", "goff");
+    // Sort to get the entry numbers of the ordered events
+    TMath::Sort(m_inputChain->GetEntries(), m_inputChain->GetV1(),
+                m_entryNumbers.data(), false);
+  }
 }
 
 ActsExamples::RootMaterialTrackReader::~RootMaterialTrackReader() {
@@ -89,11 +100,26 @@ ActsExamples::ProcessCode ActsExamples::RootMaterialTrackReader::read(
     // The collection to be written
     std::vector<Acts::RecordedMaterialTrack> mtrackCollection;
 
-    for (size_t ib = 0; ib < m_cfg.batchSize; ++ib) {
-      // Read the correct entry: batch size * event_number + ib
-      m_inputChain->GetEntry(m_cfg.batchSize * context.eventNumber + ib);
-      ACTS_VERBOSE("Reading entry: " << m_cfg.batchSize * context.eventNumber +
-                                            ib);
+    // Find the start entry and the batch size for this event
+    std::string eventNumberStr = std::to_string(context.eventNumber);
+    std::string findStartEntry = "event_id<" + eventNumberStr;
+    std::string findBatchSize = "event_id==" + eventNumberStr;
+    size_t startEntry = m_inputChain->GetEntries(findStartEntry.c_str());
+    size_t batchSize = m_inputChain->GetEntries(findBatchSize.c_str());
+    ACTS_VERBOSE("The event has " << batchSize
+                                  << " entries with the start entry "
+                                  << startEntry);
+
+    // Loop over the entries for this event
+    for (size_t ib = 0; ib < batchSize; ++ib) {
+      // Read the correct entry: startEntry + ib
+      auto entry = startEntry + ib;
+      if (not m_cfg.orderedEvents and entry < m_entryNumbers.size()) {
+        entry = m_entryNumbers[entry];
+      }
+      ACTS_VERBOSE("Reading event: " << context.eventNumber
+                                     << " with stored entry: " << entry);
+      m_inputChain->GetEntry(entry);
 
       Acts::RecordedMaterialTrack rmTrack;
       // Fill the position and momentum

Examples/Io/Root/src/RootMaterialTrackWriter.cpp

@@ -51,6 +51,7 @@ ActsExamples::RootMaterialTrackWriter::RootMaterialTrackWriter(
     throw std::bad_alloc();
 
   // Set the branches
+  m_outputTree->Branch("event_id", &m_eventId);
   m_outputTree->Branch("v_x", &m_v_x);
   m_outputTree->Branch("v_y", &m_v_y);
   m_outputTree->Branch("v_z", &m_v_z);
@@ -114,6 +115,7 @@ ActsExamples::ProcessCode ActsExamples::RootMaterialTrackWriter::writeT(
   // Exclusive access to the tree while writing
   std::lock_guard<std::mutex> lock(m_writeMutex);
 
+  m_eventId = ctx.eventNumber;
   // Loop over the material tracks and write them out
   for (auto& mtrack : materialTracks) {
     // Clearing the vector first

Examples/Io/Root/src/RootParticleReader.cpp

@@ -18,6 +18,7 @@
 
 #include <TChain.h>
 #include <TFile.h>
+#include <TMath.h>
 
 ActsExamples::RootParticleReader::RootParticleReader(
     const ActsExamples::RootParticleReader::Config& cfg)
@@ -66,6 +67,15 @@ ActsExamples::RootParticleReader::RootParticleReader(
 
   m_events = m_inputChain->GetEntries();
   ACTS_DEBUG("The full chain has " << m_events << " entries.");
+
+  // If the events are not in order, get the entry numbers for ordered events
+  if (not m_cfg.orderedEvents) {
+    m_entryNumbers.resize(m_events);
+    m_inputChain->Draw("event_id", "", "goff");
+    // Sort to get the entry numbers of the ordered events
+    TMath::Sort(m_inputChain->GetEntries(), m_inputChain->GetV1(),
+                m_entryNumbers.data(), false);
+  }
 }
 
 std::string ActsExamples::RootParticleReader::name() const {
@@ -119,9 +129,14 @@ ActsExamples::ProcessCode ActsExamples::RootParticleReader::read(
     // Secondary vertex collection
     std::vector<uint32_t> secVtxCollection;
 
-    // Read the correct entry: batch size * event_number + ib
-    m_inputChain->GetEntry(context.eventNumber);
-    ACTS_VERBOSE("Reading entry: " << context.eventNumber);
+    // Read the correct entry
+    auto entry = context.eventNumber;
+    if (not m_cfg.orderedEvents and entry < m_entryNumbers.size()) {
+      entry = m_entryNumbers[entry];
+    }
+    m_inputChain->GetEntry(entry);
+    ACTS_INFO("Reading event: " << context.eventNumber
+                                << " stored as entry: " << entry);
 
     unsigned int nParticles = m_particleId->size();
 

Examples/Io/Root/src/RootTrajectoryParametersReader.cpp

@@ -18,6 +18,7 @@
 
 #include <TChain.h>
 #include <TFile.h>
+#include <TMath.h>
 
 ActsExamples::RootTrajectoryParametersReader::RootTrajectoryParametersReader(
     const ActsExamples::RootTrajectoryParametersReader::Config& cfg)
@@ -71,6 +72,15 @@ ActsExamples::RootTrajectoryParametersReader::RootTrajectoryParametersReader(
 
   m_events = m_inputChain->GetEntries();
   ACTS_DEBUG("The full chain has " << m_events << " entries.");
+
+  // If the events are not in order, get the entry numbers for ordered events
+  if (not m_cfg.orderedEvents) {
+    m_entryNumbers.resize(m_events);
+    m_inputChain->Draw("event_nr", "", "goff");
+    // Sort to get the entry numbers of the ordered events
+    TMath::Sort(m_inputChain->GetEntries(), m_inputChain->GetV1(),
+                m_entryNumbers.data(), false);
+  }
 }
 
 std::string ActsExamples::RootTrajectoryParametersReader::name() const {
@@ -132,9 +142,14 @@ ActsExamples::ProcessCode ActsExamples::RootTrajectoryParametersReader::read(
     std::vector<Acts::BoundTrackParameters> trackParameterCollection;
     SimParticleContainer truthParticleCollection;
 
-    // Read the correct entry: batch size * event_number + ib
-    m_inputChain->GetEntry(context.eventNumber);
-    ACTS_VERBOSE("Reading entry: " << context.eventNumber);
+    // Read the correct entry
+    auto entry = context.eventNumber;
+    if (not m_cfg.orderedEvents and entry < m_entryNumbers.size()) {
+      entry = m_entryNumbers[entry];
+    }
+    m_inputChain->GetEntry(entry);
+    ACTS_INFO("Reading event: " << context.eventNumber
+                                << " stored as entry: " << entry);
 
     unsigned int nTracks = m_eLOC0_fit->size();
     for (unsigned int i = 0; i < nTracks; i++) {

Examples/Run/Vertexing/TrackReaderVertexingPerformanceWriterExample.cpp

@@ -61,7 +61,7 @@ int main(int argc, char* argv[]) {
 
   RootParticleReader::Config particleReaderConfig;
   particleReaderConfig.particleCollection = "allTruthParticles";
-  particleReaderConfig.inputFile = "particles.root";
+  particleReaderConfig.inputFile = "particles_initial.root";
   particleReaderConfig.treeName = "particles";
   particleReaderConfig.inputDir = inputDir;
   sequencer.addReader(