MCGen analysis update removing reco collisions

Author: Simple-Shyam

PWGHF/HFC/TableProducer/correlatorDplusHadrons.cxx

@@ -186,6 +186,11 @@ struct HfCorrelatorDplusHadrons {
 
   Configurable<int> selectionFlagDplus{"selectionFlagDplus", 7, "Selection Flag for Dplus"}; // 7 corresponds to topo+PID cuts
   Configurable<int> numberEventsMixed{"numberEventsMixed", 5, "Number of events mixed in ME process"};
+  Configurable<bool> removeUnreconstructedGenCollisions{"removeUnreconstructedGenCollisions",true,"Remove generator-level collisions that were not reconstructed"}; 
+  Configurable<bool> removeCollWSplitVtx{"removeCollWSplitVtx", true, "Flag for rejecting the splitted collisions"};
+  Configurable<bool> useSel8{"useSel8", true, "Flag for applying sel8 for collision selection"};
+  Configurable<bool> selNoSameBunchPileUpColl{"selNoSameBunchPileUpColl", true, "Flag for rejecting the collisions associated with the same bunch crossing"};
+  Configurable<float> zVtxMax{"zVtxMax", 10., "max. position-z of the reconstructed collision"};
   Configurable<bool> applyEfficiency{"applyEfficiency", true, "Flag for applying D-meson efficiency weights"};
   Configurable<bool> removeDaughters{"removeDaughters", true, "Flag for removing D-meson daughters from correlations"};
   Configurable<float> yCandMax{"yCandMax", 0.8, "max. cand. rapidity"};
@@ -210,6 +215,7 @@ struct HfCorrelatorDplusHadrons {
   // Event Mixing for the Data Mode
   using SelCollisionsWithDplus = soa::Filtered<soa::Join<aod::Collisions, aod::Mults, aod::EvSels, aod::DmesonSelection>>;
   using SelCollisionsWithDplusMc = soa::Filtered<soa::Join<aod::McCollisions, aod::DmesonSelection, aod::MultsExtraMC>>; // collisionFilter applied
+  using CollisionsMc = soa::Join<aod::McCollisions, aod::MultsExtraMC>;
   using CandidatesDplusData = soa::Filtered<soa::Join<aod::HfCand3Prong, aod::HfSelDplusToPiKPi, aod::HfMlDplusToPiKPi>>;
   // Event Mixing for the MCRec Mode
   using CandidatesDplusMcRec = soa::Filtered<soa::Join<aod::HfCand3Prong, aod::HfSelDplusToPiKPi, aod::HfMlDplusToPiKPi, aod::HfCand3ProngMcRec>>;
@@ -226,6 +232,8 @@ struct HfCorrelatorDplusHadrons {
   Filter dplusFilter = ((o2::aod::hf_track_index::hfflag & static_cast<uint8_t>(1 << aod::hf_cand_3prong::DecayType::DplusToPiKPi)) != static_cast<uint8_t>(0)) && aod::hf_sel_candidate_dplus::isSelDplusToPiKPi >= selectionFlagDplus;
   Filter trackFilter = (nabs(aod::track::eta) < etaTrackMax) && (nabs(aod::track::pt) > ptTrackMin) && (nabs(aod::track::dcaXY) < dcaXYTrackMax) && (nabs(aod::track::dcaZ) < dcaZTrackMax);
   Preslice<aod::McParticles> presliceMc{aod::mcparticle::mcCollisionId};
+  Preslice<CandDplusMcGen> candMcGenPerMcCollision = o2::aod::mcparticle::mcCollisionId;
+  PresliceUnsorted<soa::Join<aod::Collisions, aod::FT0Mults, aod::EvSels, aod::McCollisionLabels>> recoCollisionsPerMcCollision = o2::aod::mccollisionlabel::mcCollisionId;
   // Filter particlesFilter = nabs(aod::mcparticle::pdgCode) == 411 || ((aod::mcparticle::flags & (uint8_t)o2::aod::mcparticle::enums::PhysicalPrimary) == (uint8_t)o2::aod::mcparticle::enums::PhysicalPrimary);
   ConfigurableAxis binsMultiplicity{"binsMultiplicity", {VARIABLE_WIDTH, 0.0f, 2000.0f, 6000.0f, 100000.0f}, "Mixing bins - multiplicity"};
   ConfigurableAxis binsZVtx{"binsZVtx", {VARIABLE_WIDTH, -10.0f, -2.5f, 2.5f, 10.0f}, "Mixing bins - z-vertex"};
@@ -311,6 +319,126 @@ struct HfCorrelatorDplusHadrons {
     corrBinning = {{binsZVtx, binsMultiplicity}, true};
   }
 
+template <typename ParticleContainer, typename AssocContainer>
+void runMcGenDplusHadronAnalysis(const ParticleContainer& particlesToLoop, const AssocContainer& groupedMcParticles, int poolBin, int& counterDplusHadron)
+{
+  for (const auto& particle : particlesToLoop) {
+
+    if (std::abs(particle.pdgCode()) != Pdg::kDPlus) {
+      continue;
+    }
+
+    if (std::abs(particle.flagMcMatchGen()) != hf_decay::hf_cand_3prong::DecayChannelMain::DplusToPiKPi) {
+      continue;
+    }
+
+    double yD = RecoDecay::y(particle.pVector(), MassDPlus);
+
+    if (std::abs(yD) > yCandGenMax ||
+        particle.pt() < ptCandMin ||
+        particle.pt() > ptCandMax) {
+      continue;
+    }
+
+    std::vector<int> listDaughters{};
+    std::array<int, NDaughters> const arrDaughDplusPDG = {+kPiPlus, -kKPlus, kPiPlus};
+    std::array<int, NDaughters> prongsId{};
+
+    RecoDecay::getDaughters(particle, &listDaughters, arrDaughDplusPDG, 2);
+
+    if (listDaughters.size() != NDaughters) {
+      continue;
+    }
+
+    bool isDaughtersOk = true;
+    int counterDaughters = 0;
+
+    for (const auto& dauIdx : listDaughters) {
+      auto daughI = particlesToLoop.rawIteratorAt(dauIdx - particlesToLoop.offset());
+
+      if (std::abs(daughI.eta()) >= EtaDaughtersMax) {
+        isDaughtersOk = false;
+        break;
+      }
+
+      prongsId[counterDaughters++] = daughI.globalIndex();
+    }
+
+    if (!isDaughtersOk) {  // Skip this D+ candidate if any daughter fails eta cut
+      continue;
+    }
+    counterDplusHadron++;
+
+    registry.fill(HIST("hDplusBin"), poolBin);
+    registry.fill(HIST("hPtCandMCGen"), particle.pt());
+    registry.fill(HIST("hEtaMcGen"), particle.eta());
+    registry.fill(HIST("hPhiMcGen"), RecoDecay::constrainAngle(particle.phi(), -PIHalf));
+    registry.fill(HIST("hYMCGen"), yD);
+
+    // Prompt / Non-prompt separation
+    bool isDplusPrompt    = particle.originMcGen() == RecoDecay::OriginType::Prompt;
+    bool isDplusNonPrompt = particle.originMcGen() == RecoDecay::OriginType::NonPrompt;
+
+    if (isDplusPrompt) {
+      registry.fill(HIST("hPtCandMcGenPrompt"), particle.pt());
+    } else if (isDplusNonPrompt) {
+      registry.fill(HIST("hPtCandMcGenNonPrompt"), particle.pt());
+    }
+
+    // Count triggers
+    registry.fill(HIST("hcountDplustriggersMCGen"), 0, particle.pt());
+
+    for (const auto& particleAssoc : groupedMcParticles) {
+
+      if (std::abs(particleAssoc.eta()) > etaTrackMax ||
+          particleAssoc.pt() < ptTrackMin ||
+          particleAssoc.pt() > ptTrackMax) {
+        continue;
+      }
+
+      // Remove daughters if requested
+      if (removeDaughters) {
+        if (particleAssoc.globalIndex() == prongsId[0] ||
+            particleAssoc.globalIndex() == prongsId[1] ||
+            particleAssoc.globalIndex() == prongsId[2]) {
+          continue;
+        }
+      }
+
+      // Particle species selection
+      if ((std::abs(particleAssoc.pdgCode()) != kElectron) &&
+          (std::abs(particleAssoc.pdgCode()) != kMuonMinus) &&
+          (std::abs(particleAssoc.pdgCode()) != kPiPlus) &&
+          (std::abs(particleAssoc.pdgCode()) != kKPlus) &&
+          (std::abs(particleAssoc.pdgCode()) != kProton)) {
+        continue;
+      }
+
+      if (!particleAssoc.isPhysicalPrimary()) {
+        continue;
+      }
+
+      int trackOrigin = RecoDecay::getCharmHadronOrigin(groupedMcParticles,
+                                                        particleAssoc,
+                                                        true);
+
+      registry.fill(HIST("hPtParticleAssocMcGen"), particleAssoc.pt());
+      entryDplusHadronPair(
+        getDeltaPhi(particleAssoc.phi(), particle.phi()),
+        particleAssoc.eta() - particle.eta(),
+        particle.pt(),
+        particleAssoc.pt(),
+        poolBin
+      );
+
+      entryDplusHadronRecoInfo(MassDPlus, true);
+      entryDplusHadronGenInfo(isDplusPrompt,
+                              particleAssoc.isPhysicalPrimary(),
+                              trackOrigin);
+    }
+  }
+}
+
   /// Dplus-hadron correlation pair builder - for real data and data-like analysis (i.e. reco-level w/o matching request via MC truth)
   void processData(SelCollisionsWithDplus::iterator const& collision,
                    TracksData const& tracks,
@@ -535,103 +663,68 @@ struct HfCorrelatorDplusHadrons {
   }
   PROCESS_SWITCH(HfCorrelatorDplusHadrons, processMcRec, "Process MC Reco mode", true);
 
-  /// Dplus-Hadron correlation pair builder - for MC gen-level analysis (no filter/selection, only true signal)
-  void processMcGen(SelCollisionsWithDplusMc::iterator const& mcCollision,
-                    CandDplusMcGen const& mcParticles)
-  {
+   /// Dplus-Hadron correlation pair builder - for MC gen-level analysis (no filter/selection, only true signal)
+void processMcGen(CollisionsMc const& mcCollisions,
+                  soa::Join<aod::Collisions, aod::FT0Mults, aod::EvSels, aod::McCollisionLabels> const& collisions,
+                  CandDplusMcGen const& mcParticles)
+{
+  BinningTypeMcGen const corrBinningMcGen{{binsZVtx, binsMultiplicityMc}, true};
+
+  for (const auto& mcCollision : mcCollisions) {
+
     int counterDplusHadron = 0;
     registry.fill(HIST("hMCEvtCount"), 0);
 
-    BinningTypeMcGen const corrBinningMcGen{{binsZVtx, binsMultiplicityMc}, true};
     int poolBin = corrBinningMcGen.getBin(std::make_tuple(mcCollision.posZ(), mcCollision.multMCFT0A()));
     registry.fill(HIST("hMultFT0AMcGen"), mcCollision.multMCFT0A());
 
-    // MC gen level
-    for (const auto& particle1 : mcParticles) {
-      // check if the particle is Dplus  (for general plot filling and selection, so both cases are fine) - NOTE: decay channel is not probed!
-      if (std::abs(particle1.pdgCode()) != Pdg::kDPlus) {
-        continue;
-      }
-      if (std::abs(particle1.flagMcMatchGen()) != hf_decay::hf_cand_3prong::DecayChannelMain::DplusToPiKPi) {
-        continue;
-      }
-      double const yD = RecoDecay::y(particle1.pVector(), MassDPlus);
-      if (std::abs(yD) >= yCandMax || particle1.pt() <= ptCandMin) {
+    const auto groupedMcParticles = mcParticles.sliceBy(candMcGenPerMcCollision, mcCollision.globalIndex());
+    const auto groupedCollisions  = collisions.sliceBy(recoCollisionsPerMcCollision, mcCollision.globalIndex());
+
+    if (removeUnreconstructedGenCollisions) {
+
+      if (groupedCollisions.size() < 1) {
         continue;
       }
-      std::vector<int> listDaughters{};
-      std::array<int, NDaughters> const arrDaughDplusPDG = {+kPiPlus, -kKPlus, kPiPlus};
-      std::array<int, NDaughters> prongsId{};
-      listDaughters.clear();
-      RecoDecay::getDaughters(particle1, &listDaughters, arrDaughDplusPDG, 2);
-      int counterDaughters = 0;
-      if (listDaughters.size() != NDaughters) {
+
+      if (groupedCollisions.size() > 1 && removeCollWSplitVtx) {
         continue;
       }
-      bool isDaughtersOk = true;
-      for (const auto& dauIdx : listDaughters) {
-        auto daughI = mcParticles.rawIteratorAt(dauIdx - mcParticles.offset());
-        if (std::abs(daughI.eta()) >= EtaDaughtersMax) {
-          isDaughtersOk = false;
-          break;
-        }
-        counterDaughters += 1;
-        prongsId[counterDaughters - 1] = daughI.globalIndex();
-      }
-      if (!isDaughtersOk) {
-        continue; // Skip this D+ candidate if any daughter fails eta cut
-      }
-      counterDplusHadron++;
-
-      registry.fill(HIST("hDplusBin"), poolBin);
-      registry.fill(HIST("hPtCandMCGen"), particle1.pt());
-      registry.fill(HIST("hEtaMcGen"), particle1.eta());
-      registry.fill(HIST("hPhiMcGen"), RecoDecay::constrainAngle(particle1.phi(), -PIHalf));
-      registry.fill(HIST("hYMCGen"), yD);
 
-      // prompt and non-prompt division
-      bool isDplusPrompt = particle1.originMcGen() == RecoDecay::OriginType::Prompt;
-      bool isDplusNonPrompt = particle1.originMcGen() == RecoDecay::OriginType::NonPrompt;
-      if (isDplusPrompt) {
-        registry.fill(HIST("hPtCandMcGenPrompt"), particle1.pt());
-      } else if (isDplusNonPrompt) {
-        registry.fill(HIST("hPtCandMcGenNonPrompt"), particle1.pt());
-      }
+      for (const auto& collision : groupedCollisions) {
 
-      // Dplus Hadron correlation dedicated section
-      // if it's a Dplus particle, search for Hadron and evaluate correlations
-      registry.fill(HIST("hcountDplustriggersMCGen"), 0, particle1.pt()); // to count trigger Dplus for normalisation)
-      for (const auto& particleAssoc : mcParticles) {
-        if (std::abs(particleAssoc.eta()) > etaTrackMax || particleAssoc.pt() < ptTrackMin || particleAssoc.pt() > ptTrackMax) {
+        if (useSel8 && !collision.sel8()) {
           continue;
         }
-        if (removeDaughters) {
-          if (particleAssoc.globalIndex() == prongsId[0] || particleAssoc.globalIndex() == prongsId[1] || particleAssoc.globalIndex() == prongsId[2]) {
-            continue;
-          }
+
+        if (std::abs(collision.posZ()) > zVtxMax) {
+          continue;
         }
-        if ((std::abs(particleAssoc.pdgCode()) != kElectron) && (std::abs(particleAssoc.pdgCode()) != kMuonMinus) && (std::abs(particleAssoc.pdgCode()) != kPiPlus) && (std::abs(particleAssoc.pdgCode()) != kKPlus) && (std::abs(particleAssoc.pdgCode()) != kProton)) {
+
+        if (selNoSameBunchPileUpColl &&
+            !(collision.selection_bit(o2::aod::evsel::kNoSameBunchPileup))) {
           continue;
         }
-        if (!particleAssoc.isPhysicalPrimary()) {
+
+        if (!collision.has_mcCollision()) {
+          registry.fill(HIST("hFakeCollision"), 0.);
           continue;
         }
 
-        int trackOrigin = RecoDecay::getCharmHadronOrigin(mcParticles, particleAssoc, true);
-        registry.fill(HIST("hPtParticleAssocMcGen"), particleAssoc.pt());
-        entryDplusHadronPair(getDeltaPhi(particleAssoc.phi(), particle1.phi()),
-                             particleAssoc.eta() - particle1.eta(),
-                             particle1.pt(),
-                             particleAssoc.pt(),
-                             poolBin);
-        entryDplusHadronRecoInfo(MassDPlus, true);
-        entryDplusHadronGenInfo(isDplusPrompt, particleAssoc.isPhysicalPrimary(), trackOrigin);
-      } // end associated loop
-    } // end trigger
+       // MC particles with reconstructed-collision selection
+        runMcGenDplusHadronAnalysis(groupedMcParticles, groupedMcParticles, poolBin, counterDplusHadron);
+      }
+
+    } 
+    else {
+    // MC particles without reconstructed-collision selection (preliminay approval approach)
+      runMcGenDplusHadronAnalysis(groupedMcParticles, groupedMcParticles, poolBin, counterDplusHadron);
+    }
+
     registry.fill(HIST("hcountDplusHadronPerEvent"), counterDplusHadron);
     registry.fill(HIST("hZvtx"), mcCollision.posZ());
-    // registry.fill(HIST("hMultiplicity"), getTracksSize(mcCollision));
   }
+}
   PROCESS_SWITCH(HfCorrelatorDplusHadrons, processMcGen, "Process MC Gen mode", false);
 
   void processDataMixedEvent(SelCollisionsWithDplus const& collisions,