From 047c7f6f6aeeac2c2397c59e9491193a1615106d Mon Sep 17 00:00:00 2001
From: jparkkil <jasper.parkkila@cern.ch>
Date: Wed, 15 Apr 2026 17:37:09 +0200
Subject: [PATCH] Zvtx axis for the invariant mass. Prompt efficiency fix.

---
 PWGCF/Tasks/correlations.cxx | 22 +++++++++++-----------
 1 file changed, 11 insertions(+), 11 deletions(-)

diff --git a/PWGCF/Tasks/correlations.cxx b/PWGCF/Tasks/correlations.cxx
index 0460b80886f..8ce117c1d15 100644
--- a/PWGCF/Tasks/correlations.cxx
+++ b/PWGCF/Tasks/correlations.cxx
@@ -190,7 +190,7 @@ struct CorrelationTask {
       registry.add("etaphiTrigger", "multiplicity/centrality vs eta vs phi (triggers)", {HistType::kTH3F, {{100, 0, 100, "multiplicity/centrality"}, {100, -2, 2, "#eta"}, {200, 0, o2::constants::math::TwoPI, "#varphi"}}});
       const AxisSpec& a = AxisSpec(axisInvMass);
       AxisSpec axisSpecMass = {1000, a.binEdges[0], a.binEdges[a.getNbins()]};
-      registry.add("invMass", "2-prong invariant mass (GeV/c^2)", {HistType::kTH3F, {axisSpecMass, axisPtTrigger, axisMultiplicity}});
+      registry.add("invMass", "2-prong invariant mass (GeV/c^2)", {HistType::kTHnSparseF, {axisSpecMass, axisPtTrigger, axisMultiplicity, axisVertex}});
       if (doprocessSame2Prong2Prong || doprocessSame2Prong2ProngML) {
         registry.add("invMassTwoPart", "2D 2-prong invariant mass (GeV/c^2)", {HistType::kTHnSparseF, {axisSpecMass, axisSpecMass, axisPtTrigger, axisPtAssoc, axisMultiplicity}});
         registry.add("invMassTwoPartDPhi", "2D 2-prong invariant mass (GeV/c^2)", {HistType::kTHnSparseF, {axisSpecMass, axisSpecMass, axisPtTrigger, axisPtAssoc, axisDeltaPhi}});
@@ -370,7 +370,7 @@ struct CorrelationTask {
   using HasPDGCode = decltype(std::declval<T&>().pdgCode());
 
   template <typename TCollision, typename TTracks1, typename TTracks2>
-  void fillQA(const TCollision& collision, float multiplicity, const TTracks1& tracks1, const TTracks2& tracks2)
+  void fillQA(const TCollision& collision, float multiplicity, float posZ, const TTracks1& tracks1, const TTracks2& tracks2)
   {
     for (const auto& track1 : tracks1) {
       if constexpr (std::experimental::is_detected<HasInvMass, typename TTracks1::iterator>::value && std::experimental::is_detected<HasDecay, typename TTracks1::iterator>::value) {
@@ -380,7 +380,7 @@ struct CorrelationTask {
           if (!passMLScore(track1))
             continue;
         }
-        registry.fill(HIST("invMass"), track1.invMass(), track1.pt(), multiplicity);
+        registry.fill(HIST("invMass"), track1.invMass(), track1.pt(), multiplicity, posZ);
         for (const auto& track2 : tracks2) {
           if constexpr (std::experimental::is_detected<HasInvMass, typename TTracks2::iterator>::value && std::experimental::is_detected<HasDecay, typename TTracks2::iterator>::value) {
             if (doprocessSame2Prong2Prong || doprocessMixed2Prong2Prong || doprocessSame2Prong2ProngML || doprocessMixed2Prong2ProngML) {
@@ -874,7 +874,7 @@ struct CorrelationTask {
     registry.fill(HIST("eventcount_same"), bin);
     registry.fill(HIST("trackcount_same"), bin, tracks1.size());
     if constexpr (std::experimental::is_detected<HasDecay, typename TTracks1::iterator>::value)
-      fillQA(collision, multiplicity, tracks1, tracks2);
+      fillQA(collision, multiplicity, collision.posZ(), tracks1, tracks2);
     else
       fillQA(collision, multiplicity, tracks1);
 
@@ -1213,14 +1213,12 @@ struct CorrelationTask {
       if (std::find(cfgMcTriggerPDGs->begin(), cfgMcTriggerPDGs->end(), mcParticle.pdgCode()) != cfgMcTriggerPDGs->end()) {
         if ((mcParticle.mcDecay() != aod::cf2prongtrack::D0ToPiK) && (mcParticle.mcDecay() != aod::cf2prongtrack::D0barToKPiExclusive))
           continue; // wrong decay channel
-        if constexpr (!reflectionSpec) {
-          if (!cfgPtCentDepMLpromptSel->empty() && (mcParticle.decay() & aod::cf2prongmcpart::Prompt) == 0)
-            continue;
-        }
         if (mcParticle.cfParticleDaugh0Id() < 0 && mcParticle.cfParticleDaugh1Id() < 0)
           continue; // daughters not found
-        if constexpr (!reflectionSpec)
-          same->getTrackHistEfficiency()->Fill(CorrelationContainer::MC, mcParticle.eta(), mcParticle.pt(), 4, multiplicity, mcCollision.posZ());
+        if constexpr (!reflectionSpec) {
+          if (cfgPtCentDepMLpromptSel->empty() || (mcParticle.decay() & aod::cf2prongmcpart::Prompt) != 0)
+            same->getTrackHistEfficiency()->Fill(CorrelationContainer::MC, mcParticle.eta(), mcParticle.pt(), 4, multiplicity, mcCollision.posZ());
+        }
         p2indexCache.push_back(mcParticle.globalIndex());
       }
     }
@@ -1250,6 +1248,8 @@ struct CorrelationTask {
             return false;
           const auto& mcParticle = mcParticles.iteratorAt(*m - mcParticles.begin().globalIndex());
           if constexpr (!reflectionSpec) {
+            if (!cfgPtCentDepMLpromptSel->empty() && (mcParticle.decay() & aod::cf2prongmcpart::Prompt) == 0)
+              return true; // a valid candidate but not a prompt
             same->getTrackHistEfficiency()->Fill(CorrelationContainer::RecoPrimaries, mcParticle.eta(), mcParticle.pt(), 4, multiplicity, mcCollision.posZ());
           } else {
             if ((mcParticle.mcDecay() == aod::cf2prongtrack::D0barToKPiExclusive && (p2track.decay() == aod::cf2prongtrack::D0barToKPiExclusive || p2track.decay() == aod::cf2prongtrack::D0barToKPi)) ||
@@ -1316,7 +1316,7 @@ struct CorrelationTask {
 
     if (!(doprocessSameDerived || doprocessSameDerivedMultSet || doprocessSame2ProngDerived || doprocessSame2ProngDerivedML || doprocessSame2Prong2Prong || doprocessSame2Prong2ProngML)) {
       if constexpr (std::experimental::is_detected<HasDecay, typename Particles1::iterator>::value)
-        fillQA(mcCollision, multiplicity, mcParticles1, mcParticles2);
+        fillQA(mcCollision, multiplicity, mcCollision.posZ(), mcParticles1, mcParticles2);
       else
         fillQA(mcCollision, multiplicity, mcParticles1);
     }