Add option to fill only reconstructed MC truth collisions

Author: alicja-pp

PWGCF/FemtoUniverse/TableProducer/femtoUniverseProducerTask.cxx

@@ -135,6 +135,8 @@ struct FemtoUniverseProducerTask {
   Configurable<bool> confStoreMCmothers{"confStoreMCmothers", false, "MC truth: Fill with not only primary particles and store mothers' PDG in tempFitVar."};
   Configurable<bool> confFillCollExt{"confFillCollExt", false, "Option to fill collision extended table"};
 
+  Configurable<bool> confCollMCTruthOnlyReco{"confCollMCTruthOnlyReco", false, "Fill only MC truth collisions that were reconstructed and selected"};
+
   /// Event filtering (used for v0-cascade analysis)
   Configurable<std::string> zorroMask{"zorroMask", "", "zorro trigger class to select on (empty: none)"};
 
@@ -2601,6 +2603,7 @@ struct FemtoUniverseProducerTask {
   {
     // recos
     std::set<int> recoMcIds;
+    std::set<int> mcCollisions;
     for (const auto& col : collisions) {
       auto groupedTracks = tracks.sliceBy(perCollisionTracks, col.globalIndex());
       auto bc = col.bc_as<aod::BCsWithTimestamps>();
@@ -2616,6 +2619,7 @@ struct FemtoUniverseProducerTask {
       }
 
       if (colcheck) {
+        mcCollisions.insert(col.mcCollisionId());
         fillCollisionsCentRun3ColExtra<true>(col, ir);
         fillTracks<true>(groupedTracks);
       }
@@ -2627,6 +2631,9 @@ struct FemtoUniverseProducerTask {
 
     // truth
     for (const auto& mccol : mccols) {
+      if (confCollMCTruthOnlyReco && !mcCollisions.contains(mccol.globalIndex())) {
+        continue;
+      }
       auto groupedCollisions = collisions.sliceBy(recoCollsPerMCCollCentPbPb, mccol.globalIndex());
       for (const auto& col : groupedCollisions) {
         const auto colcheck = fillMCTruthCollisionsCentRun3(col); // fills the reco collisions for mc collision
@@ -2871,12 +2878,14 @@ struct FemtoUniverseProducerTask {
 
     // MCReco
     std::set<int> recoMcIds;
+    std::set<int> mcCollisions;
     for (const auto& col : collisions) {                                          // loop over collisions
       auto groupedTracks = tracks.sliceBy(perCollisionTracks, col.globalIndex()); // slicing for tracks
       auto groupedV0Parts = fullV0s.sliceBy(perCollisionV0s, col.globalIndex());  // slicing for V0
       getMagneticFieldTesla(col.bc_as<aod::BCsWithTimestamps>());
       const auto colcheck = fillCollisionsCentRun3<false>(col);
       if (colcheck) {
+        mcCollisions.insert(col.mcCollisionId());
         fillTracks<true>(groupedTracks);
         fillV0<true>(col, groupedV0Parts, groupedTracks);
       }
@@ -2888,6 +2897,9 @@ struct FemtoUniverseProducerTask {
 
     // MCTruth
     for (const auto& mccol : mccols) {
+      if (confCollMCTruthOnlyReco && !mcCollisions.contains(mccol.globalIndex())) {
+        continue;
+      }
       auto groupedCollisions = collisions.sliceBy(recoCollsPerMCCollCentPbPb, mccol.globalIndex()); // slicing for MC collisions
       auto groupedMCParticles = mcParticles.sliceBy(perMCCollision, mccol.globalIndex());           // slicing for MC particles
       for (const auto& col : groupedCollisions) {