From 206e5a6e235b9f9f928fe747347d18ecc5722400 Mon Sep 17 00:00:00 2001
From: Daiki Sekihata <daiki.sekihata@cern.ch>
Date: Sun, 13 Oct 2024 22:28:48 +0200
Subject: [PATCH] PWGEM/Dilepton: add protection against double counting of
 generad info

---
 PWGEM/Dilepton/Core/DileptonMC.h      | 13 ++++++++++++-
 PWGEM/Dilepton/Core/SingleTrackQCMC.h | 13 ++++++++++++-
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/PWGEM/Dilepton/Core/DileptonMC.h b/PWGEM/Dilepton/Core/DileptonMC.h
index 693e3d72552..e128115afb8 100644
--- a/PWGEM/Dilepton/Core/DileptonMC.h
+++ b/PWGEM/Dilepton/Core/DileptonMC.h
@@ -988,13 +988,22 @@ struct DileptonMC {
     // loop over mc stack and fill histograms for pure MC truth signals
     // all MC tracks which belong to the MC event corresponding to the current reconstructed event
 
+    std::vector<int> used_mccollisionIds; // used mc collisionIds
+    used_mccollisionIds.reserve(collisions.size());
+
     for (auto& collision : collisions) {
+      auto mccollision = collision.template emmcevent_as<TMCCollisions>();
+      if (std::find(used_mccollisionIds.begin(), used_mccollisionIds.end(), mccollision.globalIndex()) != used_mccollisionIds.end()) {
+        // LOGF(info, "same mc collision is repeated. continue;");
+        continue;
+      }
+      used_mccollisionIds.emplace_back(mccollision.globalIndex());
+
       float centralities[3] = {collision.centFT0M(), collision.centFT0A(), collision.centFT0C()};
       if (centralities[cfgCentEstimator] < cfgCentMin || cfgCentMax < centralities[cfgCentEstimator]) {
         continue;
       }
 
-      auto mccollision = collision.template emmcevent_as<TMCCollisions>();
       if (cfgEventGeneratorType >= 0 && mccollision.getSubGeneratorId() != cfgEventGeneratorType) {
         continue;
       }
@@ -1485,6 +1494,8 @@ struct DileptonMC {
         }
       } // end of true LS++ pair loop
     } // end of collision loop
+    used_mccollisionIds.clear();
+    used_mccollisionIds.shrink_to_fit();
   }
 
   template <typename TCollision, typename TTrack1, typename TTrack2, typename TCut>
diff --git a/PWGEM/Dilepton/Core/SingleTrackQCMC.h b/PWGEM/Dilepton/Core/SingleTrackQCMC.h
index ce2b005bba3..525383a9002 100644
--- a/PWGEM/Dilepton/Core/SingleTrackQCMC.h
+++ b/PWGEM/Dilepton/Core/SingleTrackQCMC.h
@@ -702,13 +702,22 @@ struct SingleTrackQCMC {
     // loop over mc stack and fill histograms for pure MC truth signals
     // all MC tracks which belong to the MC event corresponding to the current reconstructed event
 
+    std::vector<int> used_mccollisionIds; // used mc collisionIds
+    used_mccollisionIds.reserve(collisions.size());
+
     for (auto& collision : collisions) {
+      auto mccollision = collision.template emmcevent_as<TMCCollisions>();
+      if (std::find(used_mccollisionIds.begin(), used_mccollisionIds.end(), mccollision.globalIndex()) != used_mccollisionIds.end()) {
+        // LOGF(info, "same mc collision is repeated. continue;");
+        continue;
+      }
+      used_mccollisionIds.emplace_back(mccollision.globalIndex());
+
       float centralities[3] = {collision.centFT0M(), collision.centFT0A(), collision.centFT0C()};
       if (centralities[cfgCentEstimator] < cfgCentMin || cfgCentMax < centralities[cfgCentEstimator]) {
         continue;
       }
 
-      auto mccollision = collision.template emmcevent_as<TMCCollisions>();
       // LOGF(info, "mccollision.getGeneratorId() = %d", mccollision.getGeneratorId());
       // LOGF(info, "mccollision.getSubGeneratorId() = %d", mccollision.getSubGeneratorId());
       // LOGF(info, "mccollision.getSourceId() = %d", mccollision.getSourceId());
@@ -792,6 +801,8 @@ struct SingleTrackQCMC {
       }
 
     } // end of collision loop
+    used_mccollisionIds.clear();
+    used_mccollisionIds.shrink_to_fit();
   }
 
   std::unordered_map<int, float> map_weight; // map of track global index -> weight