AOD producer: improve MC collision labels for embedded events

In embedding scenarios where a single background event is reused across
multiple collisions, e.g.:

  Collision 0: Background 0 + Signal 0
  Collision 1: Background 0 + Signal 1
  Collision 2: Background 0 + Signal 2
  Collision 3: Background 0 + Signal 3
  Collision 4: Background 0 + Signal 4

the primary vertexer may return MC labels all pointing to "Background 0".
This caused an ambiguous lookup in the AOD producer, resulting in all
reconstructed collisions being incorrectly associated to the same
MCCollision entry.

This is addressed by collecting all MCCollision candidates matching a
given (sourceID, eventID) label and disambiguating via bunch-crossing
time, picking the MCCollision whose BC is closest to the reconstructed
vertex time.

Note: a more robust long-term solution would be to extend the primary
vertexer to return the full set of contributing MC labels per found
vertex, enabling simpler and more reliable disambiguation strategies.

This relates to https://its.cern.ch/jira/browse/O2-6840.

Author: sawenzel

Detectors/AOD/include/AODProducerWorkflow/AODProducerWorkflowSpec.h

@@ -42,6 +42,14 @@ using DataRequest = o2::globaltracking::DataRequest;
 
 namespace o2::aodproducer
 {
+/// helper struct to keep mapping of colIndex to MC labels and bunch crossing
+struct MCColInfo {
+  int colIndex;
+  int sourceID;
+  int eventID;
+  int64_t bc; // global bunch crossing
+};
+
 /// A structure or container to organize bunch crossing data of a timeframe
 /// and to facilitate fast lookup and search within bunch crossings.
 class BunchCrossings
@@ -612,7 +620,7 @@ class AODProducerWorkflowDPL : public Task
                             const gsl::span<const o2::dataformats::VtxTrackRef>& primVer2TRefs,
                             const gsl::span<const GIndex>& GIndices,
                             const o2::globaltracking::RecoContainer& data,
-                            const std::vector<std::vector<int>>& mcColToEvSrc);
+                            const std::vector<MCColInfo>& mcColToEvSrc);
 
   template <typename MCTrackLabelCursorType, typename MCMFTTrackLabelCursorType, typename MCFwdTrackLabelCursorType>
   void fillMCTrackLabelsTable(MCTrackLabelCursorType& mcTrackLabelCursor,

Detectors/AOD/src/AODProducerWorkflowSpec.cxx

@@ -910,13 +910,13 @@ void AODProducerWorkflowDPL::fillMCParticlesTable(o2::steer::MCKinematicsReader&
                                                   const gsl::span<const o2::dataformats::VtxTrackRef>& primVer2TRefs,
                                                   const gsl::span<const GIndex>& GIndices,
                                                   const o2::globaltracking::RecoContainer& data,
-                                                  const std::vector<std::vector<int>>& mcColToEvSrc)
+                                                  const std::vector<MCColInfo>& mcColToEvSrc)
 {
   int NSources = 0;
   int NEvents = 0;
   for (auto& p : mcColToEvSrc) {
-    NSources = std::max(p[1], NSources);
-    NEvents = std::max(p[2], NEvents);
+    NSources = std::max(p.sourceID, NSources);
+    NEvents = std::max(p.eventID, NEvents);
   }
   NSources++; // 0 - indexed
   NEvents++;
@@ -992,9 +992,9 @@ void AODProducerWorkflowDPL::fillMCParticlesTable(o2::steer::MCKinematicsReader&
 
   size_t offset = 0;
   for (auto& colInfo : mcColToEvSrc) { // loop over "<eventID, sourceID> <-> combined MC col. ID" key pairs
-    int event = colInfo[2];
-    int source = colInfo[1];
-    int mcColId = colInfo[0];
+    int event = colInfo.eventID;
+    int source = colInfo.sourceID;
+    int mcColId = colInfo.colIndex;
     std::vector<MCTrack> const& mcParticles = mcReader.getTracks(source, event);
     LOG(debug) << "Event=" << event << " source=" << source << " collision=" << mcColId;
     auto& preselect = mToStore[source][event];
@@ -1931,10 +1931,8 @@ void AODProducerWorkflowDPL::run(ProcessingContext& pc)
               zdcChannelsT);
   }
 
-  // keep track event/source id for each mc-collision
-  // using map and not unordered_map to ensure
-  // correct ordering when iterating over container elements
-  std::vector<std::vector<int>> mcColToEvSrc;
+  // keep track of event_id + source_id + bc for each mc-collision
+  std::vector<MCColInfo> mcColToEvSrc;
 
   if (mUseMC) {
     using namespace o2::aodmchelpers;
@@ -2007,13 +2005,13 @@ void AODProducerWorkflowDPL::run(ProcessingContext& pc)
                          0,
                          sourceID);
         }
-        mcColToEvSrc.emplace_back(std::vector<int>{iCol, sourceID, eventID}); // point background and injected signal events to one collision
+        mcColToEvSrc.emplace_back(MCColInfo{iCol, sourceID, eventID, globalBC}); // point background and injected signal events to one collision
       }
     }
   }
 
   std::sort(mcColToEvSrc.begin(), mcColToEvSrc.end(),
-            [](const std::vector<int>& left, const std::vector<int>& right) { return (left[0] < right[0]); });
+            [](const MCColInfo& left, const MCColInfo& right) { return (left.colIndex < right.colIndex); });
 
   // vector of FDD amplitudes
   int16_t aFDDAmplitudesA[8] = {0u};
@@ -2095,16 +2093,46 @@ void AODProducerWorkflowDPL::run(ProcessingContext& pc)
   }
 
   if (mUseMC) {
-    // filling MC collision labels
+    // Fill MC collision labels using information from the primary vertexer.
     mcColLabelsCursor.reserve(primVerLabels.size());
-    for (auto& label : primVerLabels) {
-      auto it = std::find_if(mcColToEvSrc.begin(), mcColToEvSrc.end(),
-                             [&label](const auto& mcColInfo) { return mcColInfo[1] == label.getSourceID() && mcColInfo[2] == label.getEventID(); });
+    for (size_t ivert = 0; ivert < primVerLabels.size(); ++ivert) {
+      const auto& label = primVerLabels[ivert];
+
+      // Collect all MC collision candidates matching this (sourceID, eventID) label.
+      // In the non-embedding case there is exactly one candidate. In the embedding
+      // case the same (sourceID, eventID) pair can appear in multiple collisions,
+      // so we need to disambiguate.
+      std::vector<std::pair<int32_t, int64_t>> candidates; // (colIndex, bc)
+      for (const auto& colInfo : mcColToEvSrc) {
+        if (colInfo.sourceID == label.getSourceID() &&
+            colInfo.eventID == label.getEventID()) {
+          candidates.emplace_back(colInfo.colIndex, colInfo.bc);
+        }
+      }
+
       int32_t mcCollisionID = -1;
-      if (it != mcColToEvSrc.end()) {
-        mcCollisionID = it->at(0);
+      if (candidates.size() == 1) {
+        mcCollisionID = candidates[0].first;
+      } else if (candidates.size() > 1) {
+        // Disambiguate by BC: pick the MCCollision whose BC is closest
+        // to the reconstructed collision's BC.
+        // TODO: Consider a complementary strategy using the MC labels of tracks
+        //       associated to the primary vertex, and/or by allowing the primary
+        //       vertexer to return multiple MC collision labels per vertex.
+        const auto& timeStamp = primVertices[ivert].getTimeStamp();
+        const double interactionTime = timeStamp.getTimeStamp() * 1E3; // us -> ns
+        const auto recoBC = relativeTime_to_GlobalBC(interactionTime);
+        int64_t bestDiff = std::numeric_limits<int64_t>::max();
+        for (const auto& [colIndex, bc] : candidates) {
+          const auto bcDiff = std::abs(static_cast<int64_t>(bc) - static_cast<int64_t>(recoBC));
+          if (bcDiff < bestDiff) {
+            bestDiff = bcDiff;
+            mcCollisionID = colIndex;
+          }
+        }
       }
-      uint16_t mcMask = 0; // todo: set mask using normalized weights?
+
+      uint16_t mcMask = 0; // TODO: set mask using normalised weights
       mcColLabelsCursor(mcCollisionID, mcMask);
     }
   }