PWGHF: Add finding leading particle in correlatorD0Hadrons for both McGen and McRec

PWGHF/HFC/TableProducer/correlatorD0Hadrons.cxx

@@ -280,6 +280,28 @@ struct HfCorrelatorD0Hadrons {
     int leadingIndex = leadingParticle.globalIndex();
     return leadingIndex;
   }
+  // ======= Find Leading Particle for McGen ============
+  template <typename TMcParticles>
+  int findLeadingParticleMcGen(TMcParticles const& mcParticles)
+  {
+    auto leadingParticle = mcParticles.begin();
+    for (auto const& mcParticle : mcParticles) {
+      if (std::abs(mcParticle.eta()) > etaTrackMax) {
+        continue;
+      }
+      if (mcParticle.pt() < ptTrackMin) {
+        continue;
+      }
+      if ((std::abs(mcParticle.pdgCode()) != kElectron) && (std::abs(mcParticle.pdgCode()) != kMuonMinus) && (std::abs(mcParticle.pdgCode()) != kPiPlus) && (std::abs(mcParticle.pdgCode()) != kKPlus) && (std::abs(mcParticle.pdgCode()) != kProton)) {
+        continue;
+      }
+      if (mcParticle.pt() > leadingParticle.pt()) {
+        leadingParticle = mcParticle;
+      }
+    }
+    int leadingIndex = leadingParticle.globalIndex();
+    return leadingIndex;
+  }
   // =======  Process starts for Data, Same event ============
 
   /// D0-h correlation pair builder - for real data and data-like analysis (i.e. reco-level w/o matching request via MC truth)
@@ -559,7 +581,7 @@ struct HfCorrelatorD0Hadrons {
           if (track.globalIndex() != leadingIndex) {
             continue;
           }
-          registry.fill(HIST("hTrackCounter"), 4); // fill no. of tracks  have leading particle
+          registry.fill(HIST("hTrackCounterRec"), 4); // fill no. of tracks  have leading particle
         }
 
         int signalStatus = 0;
@@ -605,6 +627,10 @@ struct HfCorrelatorD0Hadrons {
   {
     registry.fill(HIST("hEvtCountGen"), 0);
     // MC gen level
+    // find leading particle
+    if (correlateD0WithLeadingParticle) {
+      leadingIndex = findLeadingParticleMcGen(mcParticles);
+    }
     for (const auto& particle1 : mcParticles) {
       // check if the particle is D0 or D0bar (for general plot filling and selection, so both cases are fine) - NOTE: decay channel is not probed!
       if (std::abs(particle1.pdgCode()) != Pdg::kD0) {
@@ -647,11 +673,23 @@ struct HfCorrelatorD0Hadrons {
 
         auto indexMotherPi = RecoDecay::getMother(mcParticles, particle2, Pdg::kDStar, true, nullptr, 1); // last arguement 1 is written to consider immediate decay mother only
         auto indexMotherD0 = RecoDecay::getMother(mcParticles, particle1, Pdg::kDStar, true, nullptr, 1);
-        if (std::abs(particle2.pdgCode()) == kPiPlus && indexMotherPi >= 0 && indexMotherD0 >= 0 && indexMotherPi == indexMotherD0)
-          continue;
+        bool correlationStatus = false;
+        if (std::abs(particle2.pdgCode()) == kPiPlus && indexMotherPi >= 0 && indexMotherD0 >= 0 && indexMotherPi == indexMotherD0) {
+          if (!storeAutoCorrelationFlag) {
+            continue;
+          }
+          correlationStatus = true;
+        }
 
         registry.fill(HIST("hTrackCounterGen"), 3); // fill after soft pion removal
 
+        if (correlateD0WithLeadingParticle) {
+          if (particle2.globalIndex() != leadingIndex) {
+            continue;
+          }
+          registry.fill(HIST("hTrackCounterGen"), 4); // fill no. of tracks  have leading particle
+        }
+
         auto getTracksSize = [&mcParticles](aod::McCollision const& /*collision*/) {
           int nTracks = 0;
           for (const auto& track : mcParticles) {
@@ -664,8 +702,6 @@ struct HfCorrelatorD0Hadrons {
         using BinningTypeMcGen = FlexibleBinningPolicy<std::tuple<decltype(getTracksSize)>, aod::mccollision::PosZ, decltype(getTracksSize)>;
         BinningTypeMcGen corrBinningMcGen{{getTracksSize}, {zPoolBins, multPoolBinsMcGen}, true};
         int poolBin = corrBinningMcGen.getBin(std::make_tuple(mcCollision.posZ(), getTracksSize(mcCollision)));
-
-        bool correlationStatus = false;
         entryD0HadronPair(getDeltaPhi(particle2.phi(), particle1.phi()),
                           particle2.eta() - particle1.eta(),
                           particle1.pt(),

PWGHF/HFC/Tasks/taskCorrelationD0Hadrons.cxx

@@ -52,7 +52,7 @@ AxisSpec axisDeltaPhi = {64, -o2::constants::math::PIHalf, 3. * o2::constants::m
 AxisSpec axisPtD = {10, 0., 10., ""};
 AxisSpec axisPtHadron = {11, 0., 11., ""};
 AxisSpec axisPoolBin = {9, 0., 9., ""};
-AxisSpec axisCorrelationState = {2, 0., 1., ""};
+AxisSpec axisCorrelationState = {2, 0., 2., ""};
 ConfigurableAxis axisMass{"axisMass", {250, 1.65f, 2.15f}, ""};
 
 // definition of vectors for standard ptbin and invariant mass configurables
@@ -180,7 +180,16 @@ struct HfTaskCorrelationD0Hadrons {
      // Toward Transverse Away
      {"hToward", "Toward invmass; ptD; correlationState;entries", {HistType::kTH3F, {{axisMass}, {axisPtD}, {axisCorrelationState}}}},
      {"hTransverse", "Transverse invmass; ptD; correlationState;entries", {HistType::kTH3F, {{axisMass}, {axisPtD}, {axisCorrelationState}}}},
-     {"hAway", "Away invmass; ptD; correlationState;entries", {HistType::kTH3F, {{axisMass}, {axisPtD}, {axisCorrelationState}}}}}};
+     {"hAway", "Away invmass; ptD; correlationState;entries", {HistType::kTH3F, {{axisMass}, {axisPtD}, {axisCorrelationState}}}},
+
+     // Toward Transverse Away for McRec
+     {"hTowardRec", "Toward invmass; ptD; correlationState;entries", {HistType::kTH3F, {{axisMass}, {axisPtD}, {axisCorrelationState}}}},
+     {"hTransverseRec", "Transverse invmass; ptD; correlationState;entries", {HistType::kTH3F, {{axisMass}, {axisPtD}, {axisCorrelationState}}}},
+     {"hAwayRec", "Away invmass; ptD; correlationState;entries", {HistType::kTH3F, {{axisMass}, {axisPtD}, {axisCorrelationState}}}},
+     // Toward Transverse Away for McGen
+     {"hTowardGen", "Toward invmass; ptD; correlationState;entries", {HistType::kTH3F, {{axisMass}, {axisPtD}, {axisCorrelationState}}}},
+     {"hTransverseGen", "Transverse invmass; ptD; correlationState;entries", {HistType::kTH3F, {{axisMass}, {axisPtD}, {axisCorrelationState}}}},
+     {"hAwayGen", "Away invmass; ptD; correlationState;entries", {HistType::kTH3F, {{axisMass}, {axisPtD}, {axisCorrelationState}}}}}};
   void init(InitContext&)
   {
     int nBinsPtAxis = binsCorrelations->size() - 1;
@@ -395,12 +404,32 @@ struct HfTaskCorrelationD0Hadrons {
       int signalStatus = pairEntry.signalStatus();
       int ptBinD = o2::analysis::findBin(binsCorrelations, ptD);
       int poolBin = pairEntry.poolBin();
+      bool isAutoCorrelated = pairEntry.isAutoCorrelated();
 
       double efficiencyWeight = 1.;
       if (applyEfficiency) {
         efficiencyWeight = 1. / (efficiencyDmeson->at(o2::analysis::findBin(binsEfficiency, ptD)));
       }
-
+      if (isTowardTransverseAway) {
+        // Divide into three regions: toward, transverse, and away
+        if (ptHadron < leadingParticlePtMin) {
+          continue;
+        }
+        Region region = getRegion(deltaPhi);
+        switch (region) {
+          case Toward:
+            registry.fill(HIST("hTowardRec"), massD, ptD, isAutoCorrelated, efficiencyWeight);
+            break;
+          case Away:
+            registry.fill(HIST("hAwayRec"), massD, ptD, isAutoCorrelated, efficiencyWeight);
+            break;
+          case Transverse:
+            registry.fill(HIST("hTransverseRec"), massD, ptD, isAutoCorrelated, efficiencyWeight);
+            break;
+          default:
+            break;
+        }
+      }
       // fill correlation plots for signal/bagkground correlations
       if (pairEntry.signalStatus()) {
         registry.fill(HIST("hCorrel2DVsPtRecSig"), deltaPhi, deltaEta, ptD, ptHadron, efficiencyWeight);
@@ -589,14 +618,35 @@ struct HfTaskCorrelationD0Hadrons {
       double ptD = pairEntry.ptD();
       double ptHadron = pairEntry.ptHadron();
       int poolBin = pairEntry.poolBin();
+      double massD = pairEntry.mD();
+      bool isAutoCorrelated = pairEntry.isAutoCorrelated();
       // reject entries outside pT ranges of interest
       if (o2::analysis::findBin(binsCorrelations, ptD) < 0) {
         continue;
       }
       if (ptHadron > ptHadronMax) {
         ptHadron = ptHadronMax + 0.5;
       }
-
+      if (isTowardTransverseAway) {
+        // Divide into three regions: toward, transverse, and away
+        if (ptHadron < leadingParticlePtMin) {
+          continue;
+        }
+        Region region = getRegion(deltaPhi);
+        switch (region) {
+          case Toward:
+            registry.fill(HIST("hTowardGen"), massD, ptD, isAutoCorrelated);
+            break;
+          case Away:
+            registry.fill(HIST("hAwayGen"), massD, ptD, isAutoCorrelated);
+            break;
+          case Transverse:
+            registry.fill(HIST("hTransverseGen"), massD, ptD, isAutoCorrelated);
+            break;
+          default:
+            break;
+        }
+      }
       registry.fill(HIST("hCorrel2DVsPtGen"), deltaPhi, deltaEta, ptD, ptHadron, poolBin);
       registry.fill(HIST("hCorrel2DPtIntGen"), deltaPhi, deltaEta);
       registry.fill(HIST("hDeltaEtaPtIntGen"), deltaEta);