[PWGCF] Check MC HF decay and use dedicated efficiency bin (#11681)

Author: jaelpark

PWGCF/DataModel/CorrelationsDerived.h

@@ -11,11 +11,12 @@
 #ifndef PWGCF_DATAMODEL_CORRELATIONSDERIVED_H_
 #define PWGCF_DATAMODEL_CORRELATIONSDERIVED_H_
 
-#include <vector>
+#include "Common/DataModel/Centrality.h"
 
 #include "Framework/ASoA.h"
 #include "Framework/AnalysisDataModel.h"
-#include "Common/DataModel/Centrality.h"
+
+#include <vector>
 
 namespace o2::aod
 {
@@ -150,11 +151,16 @@ namespace cf2prongmcpart
 {
 DECLARE_SOA_INDEX_COLUMN_FULL(CFParticleDaugh0, cfParticleDaugh0, int, CFMcParticles, "_0"); //! Index to prong 1 CFMcParticle
 DECLARE_SOA_INDEX_COLUMN_FULL(CFParticleDaugh1, cfParticleDaugh1, int, CFMcParticles, "_1"); //! Index to prong 2 CFMcParticle
+DECLARE_SOA_COLUMN(Decay, decay, uint8_t);                                                   //! Particle decay and flags
+enum ParticleDecayFlags {
+  Prompt = 0x80
+};
 } // namespace cf2prongmcpart
 DECLARE_SOA_TABLE(CF2ProngMcParts, "AOD", "CF2PRONGMCPART", //! Table for the daughter particles of a 2-prong particle, to be joined with CFMcParticles
                   o2::soa::Index<>,
                   cf2prongmcpart::CFParticleDaugh0Id,
-                  cf2prongmcpart::CFParticleDaugh1Id)
+                  cf2prongmcpart::CFParticleDaugh1Id,
+                  cf2prongmcpart::Decay)
 using CF2ProngMcPart = CF2ProngMcParts::iterator;
 
 } // namespace o2::aod

PWGCF/TableProducer/filter2Prong.cxx

@@ -141,13 +141,14 @@ struct Filter2Prong {
   }
   PROCESS_SWITCH(Filter2Prong, processData, "Process data D0 candidates", true);
 
-  void processMC(aod::McCollisions::iterator const&, aod::CFMcParticleRefs const& cfmcparticles, [[maybe_unused]] aod::McParticles const& mcparticles)
+  using HFMCTrack = soa::Join<aod::McParticles, aod::HfCand2ProngMcGen>;
+  void processMC(aod::McCollisions::iterator const&, aod::CFMcParticleRefs const& cfmcparticles, [[maybe_unused]] HFMCTrack const& mcparticles)
   {
     // The main filter outputs the primary MC particles. Here we just resolve the daughter indices that are needed for the efficiency matching.
     for (const auto& r : cfmcparticles) {
-      const auto& mcParticle = r.mcParticle();
-      if (mcParticle.daughtersIds().size() != 2) {
-        output2ProngMcParts(-1, -1);
+      const auto& mcParticle = r.mcParticle_as<HFMCTrack>();
+      if ((mcParticle.flagMcMatchGen() & (1 << aod::hf_cand_2prong::DecayType::D0ToPiK)) == 0 || mcParticle.daughtersIds().size() != 2) {
+        output2ProngMcParts(-1, -1, aod::cf2prongtrack::Generic2Prong);
         continue;
       }
       int prongCFId[2] = {-1, -1};
@@ -159,7 +160,8 @@ struct Filter2Prong {
           }
         }
       }
-      output2ProngMcParts(prongCFId[0], prongCFId[1]);
+      output2ProngMcParts(prongCFId[0], prongCFId[1],
+                          (mcParticle.pdgCode() >= 0 ? aod::cf2prongtrack::D0ToPiK : aod::cf2prongtrack::D0barToKPi) | ((mcParticle.originMcGen() & RecoDecay::OriginType::Prompt) ? aod::cf2prongmcpart::Prompt : 0));
     }
   }
   PROCESS_SWITCH(Filter2Prong, processMC, "Process MC 2-prong daughters", false);

PWGCF/Tasks/correlations.cxx

@@ -987,9 +987,11 @@ struct CorrelationTask {
       case 2212: // proton
       case -2212:
         return 2;
-      default: // NOTE. The efficiency histogram is hardcoded to contain 4 species. Anything special will have the last slot.
-        return 3;
     }
+    if (std::find(cfgMcTriggerPDGs->begin(), cfgMcTriggerPDGs->end(), pdgCode) != cfgMcTriggerPDGs->end())
+      return 4; // NOTE - if changed, the number in processMCEfficiency2Prong needs to be changed too since we skip the getSpecies call
+    else        // The efficiency histogram is hardcoded to contain 5 species. Anything special will have the 4th slot.
+      return 3;
   }
 
   // NOTE SmallGroups includes soa::Filtered always
@@ -1054,10 +1056,12 @@ struct CorrelationTask {
     // Primaries
     p2indexCache.clear();
     for (const auto& mcParticle : mcParticles) {
-      if (mcParticle.isPhysicalPrimary() && std::find(cfgMcTriggerPDGs->begin(), cfgMcTriggerPDGs->end(), mcParticle.pdgCode()) != cfgMcTriggerPDGs->end()) {
-        same->getTrackHistEfficiency()->Fill(CorrelationContainer::MC, mcParticle.eta(), mcParticle.pt(), getSpecies(mcParticle.pdgCode()), multiplicity, mcCollision.posZ());
+      if (std::find(cfgMcTriggerPDGs->begin(), cfgMcTriggerPDGs->end(), mcParticle.pdgCode()) != cfgMcTriggerPDGs->end()) {
+        if (mcParticle.decay() != aod::cf2prongtrack::D0ToPiK || mcParticle.decay() != aod::cf2prongtrack::D0barToKPi)
+          continue; // wrong decay channel
+        same->getTrackHistEfficiency()->Fill(CorrelationContainer::MC, mcParticle.eta(), mcParticle.pt(), 4, multiplicity, mcCollision.posZ());
         if (mcParticle.cfParticleDaugh0Id() < 0 || mcParticle.cfParticleDaugh1Id() < 0)
-          continue;
+          continue; // daughters not found
         p2indexCache.push_back(mcParticle.globalIndex());
       }
     }
@@ -1068,26 +1072,25 @@ struct CorrelationTask {
         // Check if the mc particles of the prongs are found.
         if (p2track.cfTrackProng0Id() < 0 || p2track.cfTrackProng1Id() < 0) {
           // fake track
-          same->getTrackHistEfficiency()->Fill(CorrelationContainer::Fake, p2track.eta(), p2track.pt(), 0, multiplicity, mcCollision.posZ());
+          same->getTrackHistEfficiency()->Fill(CorrelationContainer::Fake, p2track.eta(), p2track.pt(), 4, multiplicity, mcCollision.posZ());
+          continue;
         }
         const auto& p0 = p2track.cfTrackProng0_as<aod::CFTracksWithLabel>();
         const auto& p1 = p2track.cfTrackProng1_as<aod::CFTracksWithLabel>();
         if (p0.has_cfMCParticle() && p1.has_cfMCParticle()) {
           // find the 2-prong MC particle by the daughter MC particle IDs
           auto m = std::find_if(p2indexCache.begin(), p2indexCache.end(), [&](const auto& t) -> bool {
             const auto& mcParticle = mcParticles.iteratorAt(t - mcParticles.begin().globalIndex());
-            return p0.cfMCParticleId() == mcParticle.cfParticleDaugh0Id() && p1.cfMCParticleId() == mcParticle.cfParticleDaugh1Id();
+            return (p0.cfMCParticleId() == mcParticle.cfParticleDaugh0Id() && p1.cfMCParticleId() == mcParticle.cfParticleDaugh1Id()) || (p0.cfMCParticleId() == mcParticle.cfParticleDaugh1Id() && p1.cfMCParticleId() == mcParticle.cfParticleDaugh0Id());
           });
           if (m == p2indexCache.end())
             continue;
           const auto& mcParticle = mcParticles.iteratorAt(*m - mcParticles.begin().globalIndex());
-          if (mcParticle.isPhysicalPrimary()) {
-            same->getTrackHistEfficiency()->Fill(CorrelationContainer::RecoPrimaries, mcParticle.eta(), mcParticle.pt(), getSpecies(mcParticle.pdgCode()), multiplicity, mcCollision.posZ());
-          }
-          same->getTrackHistEfficiency()->Fill(CorrelationContainer::RecoAll, mcParticle.eta(), mcParticle.pt(), getSpecies(mcParticle.pdgCode()), multiplicity, mcCollision.posZ());
+          same->getTrackHistEfficiency()->Fill(CorrelationContainer::RecoPrimaries, mcParticle.eta(), mcParticle.pt(), 4, multiplicity, mcCollision.posZ());
+          same->getTrackHistEfficiency()->Fill(CorrelationContainer::RecoAll, mcParticle.eta(), mcParticle.pt(), 4, multiplicity, mcCollision.posZ());
         } else {
           // fake track
-          same->getTrackHistEfficiency()->Fill(CorrelationContainer::Fake, p2track.eta(), p2track.pt(), 0, multiplicity, mcCollision.posZ());
+          same->getTrackHistEfficiency()->Fill(CorrelationContainer::Fake, p2track.eta(), p2track.pt(), 4, multiplicity, mcCollision.posZ());
         }
       }
     }