[PWGHF/HFC] Adding initial setup for correlation and event mixing at MC truth level

PWGHF/HFC/DataModel/CorrelationTables.h

@@ -457,6 +457,10 @@ DECLARE_SOA_COLUMN(PoolBin, poolBin, int);
 DECLARE_SOA_DYNAMIC_COLUMN(DeltaEta, deltaEta, [](float etaTrack, float etaCandidate) -> float { return (etaTrack - etaCandidate); });
 DECLARE_SOA_DYNAMIC_COLUMN(DeltaPhi, deltaPhi, [](float phiCandidate, float phiTrack) -> float { return RecoDecay::constrainAngle(phiTrack - phiCandidate, -o2::constants::math::PIHalf); });
 DECLARE_SOA_DYNAMIC_COLUMN(DeltaM, deltaM, [](float massDstar, float massD0) -> float { return (massDstar - massD0); });
+// MC Gen info columns
+DECLARE_SOA_COLUMN(IsPrompt, isPrompt, bool);                   //! Used in MC-Gen, D* Prompt or Non-Prompt
+DECLARE_SOA_COLUMN(IsPhysicalPrimary, isPhysicalPrimary, bool); //! Used in MC-Gen, primary associated particle flag
+DECLARE_SOA_COLUMN(TrackOrigin, trackOrigin, int);              //! Used in MC-Gen, associated particle charm hadron origin
 } // namespace hf_correlation_dstar_hadron
 
 DECLARE_SOA_TABLE(DstarHadronPair, "AOD", "DSTRHPAIR", // D* Hadrons pairs Informations
@@ -493,6 +497,22 @@ DECLARE_SOA_TABLE(Dstar, "AOD", "DSTAR", // Only Dstar properties
                   hf_correlation_dstar_hadron::TimeStamp,
                   hf_correlation_dstar_hadron::PoolBin);
 
+DECLARE_SOA_TABLE(DstarHadronMcGenPair, "AOD", "DSTRHGENP", //! D*-Hadron MC Gen pairs
+                  hf_correlation_dstar_hadron::PhiDstar,
+                  hf_correlation_dstar_hadron::EtaDstar,
+                  hf_correlation_dstar_hadron::PtDstar,
+                  hf_correlation_dstar_hadron::PhiTrack,
+                  hf_correlation_dstar_hadron::EtaTrack,
+                  hf_correlation_dstar_hadron::PtTrack,
+                  hf_correlation_dstar_hadron::PoolBin,
+                  hf_correlation_dstar_hadron::DeltaPhi<hf_correlation_dstar_hadron::PhiDstar, hf_correlation_dstar_hadron::PhiTrack>,
+                  hf_correlation_dstar_hadron::DeltaEta<hf_correlation_dstar_hadron::EtaDstar, hf_correlation_dstar_hadron::EtaTrack>);
+
+DECLARE_SOA_TABLE(DstarHadronGenInfo, "AOD", "DSTRHGENINFO", //! D*-Hadron MC Gen information
+                  hf_correlation_dstar_hadron::IsPrompt,
+                  hf_correlation_dstar_hadron::IsPhysicalPrimary,
+                  hf_correlation_dstar_hadron::TrackOrigin);
+
 // Note: Table for selection of Lc in a collision
 namespace hf_selection_lc_collision
 {

PWGHF/HFC/TableProducer/correlatorDstarHadrons.cxx

@@ -15,12 +15,14 @@
 
 /// \brief Correlator for D* and hadrons. This task is used to produce table for D* and hadron pairs.
 
+#include "PWGHF/Core/DecayChannels.h"
 #include "PWGHF/DataModel/AliasTables.h"
 #include "PWGHF/DataModel/CandidateReconstructionTables.h"
 #include "PWGHF/DataModel/CandidateSelectionTables.h"
 #include "PWGHF/HFC/DataModel/CorrelationTables.h"
 #include "PWGHF/Utils/utilsAnalysis.h"
 
+#include "Common/Core/RecoDecay.h"
 #include "Common/DataModel/Multiplicity.h"
 #include "Common/DataModel/TrackSelectionTables.h"
 
@@ -40,6 +42,8 @@
 #include <Framework/Logger.h>
 #include <Framework/runDataProcessing.h>
 
+#include <TPDGCode.h>
+
 #include <array>
 #include <cstdlib>
 #include <numeric>
@@ -50,6 +54,8 @@ using namespace o2;
 using namespace o2::framework;
 using namespace o2::framework::expressions;
 
+using BinningTypeMcGen = ColumnBinningPolicy<aod::mccollision::PosZ, o2::aod::mult::MultMCFT0A>;
+
 const int nBinsPtCorrelation = 8;
 
 const double binsPtCorrelationsDefault[nBinsPtCorrelation + 1] = {0., 2., 4., 6., 8., 12., 16., 24., 100.};
@@ -118,6 +124,8 @@ struct HfCorrelatorDstarHadrons {
   Produces<aod::DstarHadronPair> rowsDstarHadronPair;
   Produces<aod::Dstar> rowsDstar;
   Produces<aod::Hadron> rowsAssoTrack;
+  Produces<aod::DstarHadronMcGenPair> rowsDstarHadronMcGenPair;
+  Produces<aod::DstarHadronGenInfo> rowsDstarHadronGenInfo;
 
   // Enable separate tables for Dstar and Track for offline Event mixing
   Configurable<bool> enableSeparateTables{"enableSeparateTables", false, "Enable separate tables for Dstar and Track for offline Event mixing"};
@@ -137,6 +145,7 @@ struct HfCorrelatorDstarHadrons {
   Configurable<float> dcazAssoTrackMax{"dcazAssoTrackMax", 10.0, "max DCAz of Associated Track"};
   Configurable<float> ptAssoTrackMin{"ptAssoTrackMin", 0.5, "min Pt of Associated Track"};
   Configurable<float> ptAssoTrackMax{"ptAssoTrackMax", 50.0, "max pT of Associated Track"};
+  Configurable<int> numberEventsMixed{"numberEventsMixed", 5, "number of events to mix"};
 
   Configurable<std::vector<double>> binsPtCorrelations{"binsPtCorrelations", std::vector<double>{vecBinsPtCorrelationsDefault}, "pT bin limits for correlation plots"};
   Configurable<std::vector<double>> signalRegionLefBound{"signalRegionLefBound", std::vector<double>{vecSignalRegionLefBoundDefault}, "left boundary of signal region vs pT"};
@@ -163,6 +172,11 @@ struct HfCorrelatorDstarHadrons {
 
   using FilteredTracks = soa::Filtered<aod::TracksWDca>;
 
+  // MC Gen type aliases
+  using CandDstarMcGen = soa::Join<aod::McParticles, aod::HfCandDstarMcGen>;
+  using McCollisionsWithMult = soa::Join<aod::McCollisions, aod::MultsExtraMC>;
+  using RecoCollisionsForMcGen = soa::Join<aod::Collisions, aod::McCollisionLabels>;
+
   // collision table filter
   Filter collisionFilter = aod::hf_selection_dmeson_collision::dmesonSel == selectOnlyCollisionWDstar;
   // candidate filter
@@ -177,6 +191,10 @@ struct HfCorrelatorDstarHadrons {
   // Preslice<aod::TracksWDca> perColTracks = aod::track::collisionId;
   Preslice<FilteredTracks> perColTracks = aod::track::collisionId;
 
+  // MC Gen preslices
+  Preslice<CandDstarMcGen> perCollisionCandMcGen = o2::aod::mcparticle::mcCollisionId;
+  PresliceUnsorted<RecoCollisionsForMcGen> collPerMcCollision = o2::aod::mccollisionlabel::mcCollisionId;
+
   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"};
   BinningType binningScheme{{binsZVtx, binsMultiplicity}, true};
@@ -192,7 +210,7 @@ struct HfCorrelatorDstarHadrons {
 
   void init(InitContext&)
   {
-    std::array<bool, 2> processes = {doprocessDataSameEvent, doprocessDataWithMixedEvent};
+    std::array<bool, 4> processes = {doprocessDataSameEvent, doprocessDataWithMixedEvent, doprocessSeMcGen, doprocessMcGenME};
     if (std::accumulate(processes.begin(), processes.end(), 0) != 1) {
       LOGP(fatal, "One and only one process function must be enabled at a time.");
     }
@@ -406,6 +424,140 @@ struct HfCorrelatorDstarHadrons {
 
   } // processDataWithMixedEvent
   PROCESS_SWITCH(HfCorrelatorDstarHadrons, processDataWithMixedEvent, "process only mixed events data", false);
+
+  /// D*-Hadron correlation pair builder at MC Gen same-event level (true signal only, no mass selection)
+  void processSeMcGen(McCollisionsWithMult const& mcCollisions,
+                      RecoCollisionsForMcGen const& collisions,
+                      CandDstarMcGen const& mcParticles)
+  {
+    BinningTypeMcGen const corrBinningMcGen{{binsZVtx, binsMultiplicity}, true};
+
+    for (const auto& mcCollision : mcCollisions) {
+      int const poolBin = corrBinningMcGen.getBin(std::make_tuple(mcCollision.posZ(), mcCollision.multMCFT0A()));
+
+      const auto groupedMcParticles = mcParticles.sliceBy(perCollisionCandMcGen, mcCollision.globalIndex());
+      const auto groupedCollisions = collisions.sliceBy(collPerMcCollision, mcCollision.globalIndex());
+
+      if (groupedCollisions.size() < 1) {
+        continue; // skip MC events with no reconstructed collision
+      }
+
+      // loop over gen-level D* candidates
+      for (const auto& particle : groupedMcParticles) {
+        if (std::abs(particle.flagMcMatchGen()) != hf_decay::hf_cand_dstar::DecayChannelMain::DstarToPiKPi) {
+          continue;
+        }
+        auto yDstar = RecoDecay::y(particle.pVector(), constants::physics::MassDStar);
+        if (std::abs(yDstar) > yAbsDstarMax) {
+          continue;
+        }
+        bool const isDstarPrompt = particle.originMcGen() == RecoDecay::OriginType::Prompt;
+
+        // find D* daughters (K, pi from D0, soft pi) to exclude from associated particles
+        std::vector<int> listDaughters{};
+        std::array<int, 3> const arrDaughDstarPDG = {kKPlus, kPiPlus, kPiPlus};
+        RecoDecay::getDaughters(particle, &listDaughters, arrDaughDstarPDG, 2);
+
+        // loop over associated particles (charged hadrons at gen level)
+        for (const auto& particleAssoc : groupedMcParticles) {
+          if (std::abs(particleAssoc.eta()) > etaAbsAssoTrackMax ||
+              particleAssoc.pt() < ptAssoTrackMin ||
+              particleAssoc.pt() > ptAssoTrackMax) {
+            continue;
+          }
+          if (!particleAssoc.isPhysicalPrimary()) {
+            continue;
+          }
+          // select only charged hadrons: e, mu, pi, K, p
+          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;
+          }
+          // skip D* daughter particles
+          bool isDaughter = false;
+          for (const auto& dauIdx : listDaughters) {
+            if (particleAssoc.globalIndex() == dauIdx) {
+              isDaughter = true;
+              break;
+            }
+          }
+          if (isDaughter) {
+            continue;
+          }
+
+          int const trackOrigin = RecoDecay::getCharmHadronOrigin(groupedMcParticles, particleAssoc, true);
+
+          rowsDstarHadronMcGenPair(particle.phi(),
+                                   particle.eta(),
+                                   particle.pt(),
+                                   particleAssoc.phi(),
+                                   particleAssoc.eta(),
+                                   particleAssoc.pt(),
+                                   poolBin);
+          rowsDstarHadronGenInfo(isDstarPrompt,
+                                 particleAssoc.isPhysicalPrimary(),
+                                 trackOrigin);
+        } // associated particle loop
+      } // D* gen particle loop
+    } // McCollision loop
+  } // processSeMcGen
+  PROCESS_SWITCH(HfCorrelatorDstarHadrons, processSeMcGen, "Process MC Gen same-event mode", false);
+
+  /// D*-Hadron correlation pair builder at MC Gen mixed-event level
+  void processMcGenME(McCollisionsWithMult const& collisions,
+                      CandDstarMcGen const& mcParticles)
+  {
+    BinningTypeMcGen const corrBinningMcGen{{binsZVtx, binsMultiplicity}, true};
+    auto tracksTuple = std::make_tuple(mcParticles, mcParticles);
+    Pair<McCollisionsWithMult, CandDstarMcGen, CandDstarMcGen, BinningTypeMcGen> const pairMcGen{corrBinningMcGen, numberEventsMixed, -1, collisions, tracksTuple, &cache};
+
+    for (const auto& [c1, tracks1, c2, tracks2] : pairMcGen) {
+      int const poolBin = corrBinningMcGen.getBin(std::make_tuple(c1.posZ(), c1.multMCFT0A()));
+      for (const auto& [particle, particleAssoc] : o2::soa::combinations(o2::soa::CombinationsFullIndexPolicy(tracks1, tracks2))) {
+        if (std::abs(particle.flagMcMatchGen()) != hf_decay::hf_cand_dstar::DecayChannelMain::DstarToPiKPi) {
+          continue;
+        }
+        auto yDstar = RecoDecay::y(particle.pVector(), constants::physics::MassDStar);
+        if (std::abs(yDstar) > yAbsDstarMax) {
+          continue;
+        }
+        if (std::abs(particleAssoc.eta()) > etaAbsAssoTrackMax ||
+            particleAssoc.pt() < ptAssoTrackMin ||
+            particleAssoc.pt() > ptAssoTrackMax) {
+          continue;
+        }
+        if (!particleAssoc.isPhysicalPrimary()) {
+          continue;
+        }
+        // select only charged hadrons: e, mu, pi, K, p
+        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;
+        }
+        bool const isDstarPrompt = particle.originMcGen() == RecoDecay::OriginType::Prompt;
+        // use full mcParticles (not sliced) since particles come from different collisions in ME
+        int const trackOrigin = RecoDecay::getCharmHadronOrigin(mcParticles, particleAssoc, true);
+
+        rowsDstarHadronMcGenPair(particle.phi(),
+                                 particle.eta(),
+                                 particle.pt(),
+                                 particleAssoc.phi(),
+                                 particleAssoc.eta(),
+                                 particleAssoc.pt(),
+                                 poolBin);
+        rowsDstarHadronGenInfo(isDstarPrompt,
+                               particleAssoc.isPhysicalPrimary(),
+                               trackOrigin);
+      } // associated particle loop
+    } // ME pair loop
+  } // processMcGenME
+  PROCESS_SWITCH(HfCorrelatorDstarHadrons, processMcGenME, "Process MC Gen mixed-event mode", false);
 };
 
 WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)

PWGHF/HFC/Tasks/taskCorrelationDstarHadrons.cxx

@@ -14,6 +14,8 @@
 /// \author Fabrizio Grosa <fabrizio.grosa@cern.ch>, CERN
 /// \author Shyam Kumar <shyam.kumar@cern.ch>
 
+/// \brief Task to strore correlations between D* and hadrons.
+
 #include "PWGHF/Core/SelectorCuts.h"
 #include "PWGHF/HFC/DataModel/CorrelationTables.h"
 #include "PWGHF/Utils/utilsAnalysis.h"
@@ -131,14 +133,22 @@ struct HfTaskCorrelationDstarHadrons {
     AxisSpec const axisSpecPtHadron = {ptHadronBinsEdges};
     AxisSpec const axisSpecPoolBin = {9, 0., 9.};
 
-    registry.add("hCorrel2DVsPtSignalRegion", stringDHadron + stringSignal + stringDeltaPhi + stringDeltaEta + stringPtD + stringPtHadron + stringPoolBin + "entries", {HistType::kTHnSparseD, {axisSpecDeltaPhi, axisSpecDeltaEta, axisSpecPtDstar, axisSpecPtHadron, axisSpecPoolBin}}, true);
-    registry.add("hCorrel2DPtIntSignalRegion", stringDHadron + stringSignal + stringDeltaPhi + stringDeltaEta + "entries", {HistType::kTH2D, {axisSpecDeltaPhi, axisSpecDeltaEta}}, true);
-    registry.add("hDeltaEtaPtIntSignalRegion", stringDHadron + stringSignal + stringDeltaEta + "entries", {HistType::kTH1D, {axisSpecDeltaEta}}, true);
-    registry.add("hDeltaPhiPtIntSignalRegion", stringDHadron + stringSignal + stringDeltaPhi + "entries", {HistType::kTH1D, {axisSpecDeltaPhi}}, true);
-    registry.add("hCorrel2DVsPtSidebands", stringDHadron + stringSideband + stringDeltaPhi + stringDeltaEta + stringPtD + stringPtHadron + stringPoolBin + "entries", {HistType::kTHnSparseD, {axisSpecDeltaPhi, axisSpecDeltaEta, axisSpecPtDstar, axisSpecPtHadron, axisSpecPoolBin}}, true);
-    registry.add("hCorrel2DPtIntSidebands", stringDHadron + stringSideband + stringDeltaPhi + stringDeltaEta + "entries", {HistType::kTH2D, {axisSpecDeltaPhi, axisSpecDeltaEta}}, true);
-    registry.add("hDeltaEtaPtIntSidebands", stringDHadron + stringSideband + stringDeltaEta + "entries", {HistType::kTH1D, {axisSpecDeltaEta}}, true);
-    registry.add("hDeltaPhiPtIntSidebands", stringDHadron + stringSideband + stringDeltaPhi + "entries", {HistType::kTH1D, {axisSpecDeltaPhi}}, true);
+    if (doprocessData) {
+      registry.add("hCorrel2DVsPtSignalRegion", stringDHadron + stringSignal + stringDeltaPhi + stringDeltaEta + stringPtD + stringPtHadron + stringPoolBin + "entries", {HistType::kTHnSparseD, {axisSpecDeltaPhi, axisSpecDeltaEta, axisSpecPtDstar, axisSpecPtHadron, axisSpecPoolBin}}, true);
+      registry.add("hCorrel2DPtIntSignalRegion", stringDHadron + stringSignal + stringDeltaPhi + stringDeltaEta + "entries", {HistType::kTH2D, {axisSpecDeltaPhi, axisSpecDeltaEta}}, true);
+      registry.add("hDeltaEtaPtIntSignalRegion", stringDHadron + stringSignal + stringDeltaEta + "entries", {HistType::kTH1D, {axisSpecDeltaEta}}, true);
+      registry.add("hDeltaPhiPtIntSignalRegion", stringDHadron + stringSignal + stringDeltaPhi + "entries", {HistType::kTH1D, {axisSpecDeltaPhi}}, true);
+      registry.add("hCorrel2DVsPtSidebands", stringDHadron + stringSideband + stringDeltaPhi + stringDeltaEta + stringPtD + stringPtHadron + stringPoolBin + "entries", {HistType::kTHnSparseD, {axisSpecDeltaPhi, axisSpecDeltaEta, axisSpecPtDstar, axisSpecPtHadron, axisSpecPoolBin}}, true);
+      registry.add("hCorrel2DPtIntSidebands", stringDHadron + stringSideband + stringDeltaPhi + stringDeltaEta + "entries", {HistType::kTH2D, {axisSpecDeltaPhi, axisSpecDeltaEta}}, true);
+      registry.add("hDeltaEtaPtIntSidebands", stringDHadron + stringSideband + stringDeltaEta + "entries", {HistType::kTH1D, {axisSpecDeltaEta}}, true);
+      registry.add("hDeltaPhiPtIntSidebands", stringDHadron + stringSideband + stringDeltaPhi + "entries", {HistType::kTH1D, {axisSpecDeltaPhi}}, true);
+    }
+    // MC Gen histograms
+    if (doprocessMcGen) {
+      registry.add("hCorrel2DVsPtMcGen", stringDHadron + "MC Gen;" + stringDeltaPhi + stringDeltaEta + stringPtD + stringPtHadron + stringPoolBin + "entries", {HistType::kTHnSparseD, {axisSpecDeltaPhi, axisSpecDeltaEta, axisSpecPtDstar, axisSpecPtHadron, axisSpecPoolBin}}, true);
+      registry.add("hCorrel2DVsPtMcGenPrompt", stringDHadron + "MC Gen Prompt;" + stringDeltaPhi + stringDeltaEta + stringPtD + stringPtHadron + stringPoolBin + "entries", {HistType::kTHnSparseD, {axisSpecDeltaPhi, axisSpecDeltaEta, axisSpecPtDstar, axisSpecPtHadron, axisSpecPoolBin}}, true);
+      registry.add("hCorrel2DVsPtMcGenNonPrompt", stringDHadron + "MC Gen NonPrompt;" + stringDeltaPhi + stringDeltaEta + stringPtD + stringPtHadron + stringPoolBin + "entries", {HistType::kTHnSparseD, {axisSpecDeltaPhi, axisSpecDeltaEta, axisSpecPtDstar, axisSpecPtHadron, axisSpecPoolBin}}, true);
+    }
 
     if (applyEfficiency && useCcdbEfficiency) {
       ccdbApi.init(ccdbUrl);
@@ -224,7 +234,7 @@ struct HfTaskCorrelationDstarHadrons {
           efficiencyWeightDstar = vecHistEfficiencyDstar[0]->GetBinContent(1);
         } else {
           efficiencyWeightDstar = vecHistEfficiencyDstar[0]->GetBinContent(vecHistEfficiencyDstar[0]->GetXaxis()->FindBin(ptDstar));
-          if (!efficiencyWeightDstar) {
+          if (efficiencyWeightDstar == 0.0) {
             LOGF(fatal, "Dstar efficiency weight can't be zero.");
           }
         }
@@ -233,7 +243,7 @@ struct HfTaskCorrelationDstarHadrons {
           efficiencyWeightTracks = vecHistEfficiencyTracks[0]->GetBinContent(1);
         } else {
           efficiencyWeightTracks = vecHistEfficiencyTracks[0]->GetBinContent(vecHistEfficiencyTracks[0]->GetXaxis()->FindBin(ptTrack));
-          if (!efficiencyWeightTracks) {
+          if (efficiencyWeightTracks == 0.0) {
             LOGF(fatal, "track efficiency weight can't be zero");
           }
         }
@@ -259,6 +269,27 @@ struct HfTaskCorrelationDstarHadrons {
     }
   }
   PROCESS_SWITCH(HfTaskCorrelationDstarHadrons, processData, " process data only", true);
+
+  /// Fill THnSparse histograms with D*-hadron pairs at MC Gen level (both SE and ME)
+  void processMcGen(soa::Join<aod::DstarHadronMcGenPair, aod::DstarHadronGenInfo> const& pairsMcGen)
+  {
+    for (const auto& pair : pairsMcGen) {
+      float const deltaPhi = pair.deltaPhi();
+      float const deltaEta = pair.deltaEta();
+      float const ptDstar = pair.ptDstar();
+      float const ptTrack = pair.ptTrack();
+      int const poolBin = pair.poolBin();
+      bool const isPrompt = pair.isPrompt();
+
+      registry.fill(HIST("hCorrel2DVsPtMcGen"), deltaPhi, deltaEta, ptDstar, ptTrack, poolBin);
+      if (isPrompt) {
+        registry.fill(HIST("hCorrel2DVsPtMcGenPrompt"), deltaPhi, deltaEta, ptDstar, ptTrack, poolBin);
+      } else {
+        registry.fill(HIST("hCorrel2DVsPtMcGenNonPrompt"), deltaPhi, deltaEta, ptDstar, ptTrack, poolBin);
+      }
+    }
+  }
+  PROCESS_SWITCH(HfTaskCorrelationDstarHadrons, processMcGen, "Process MC Gen", false);
 };
 
 WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)