[PWGHF] Add identified particle selection in efficiency process (#14161)

Author: scattaru

PWGHF/HFC/TableProducer/correlatorDsHadrons.cxx

@@ -598,24 +598,41 @@ struct HfCorrelatorDsHadrons {
         if (!track.isGlobalTrackWoDCA()) {
           continue;
         }
+        // apply PID selection
+        if (pidTrkApplied) {
+          if (!passPIDSelection(track, trkPIDspecies, pidTPCMax, pidTOFMax, tofPIDThreshold, forceTOF)) {
+            continue;
+          }
+          registry.fill(HIST("hTpcTofNSigmaPidPionMc"), o2::aod::pidutils::tpcNSigma(o2::track::PID::Pion, track), o2::aod::pidutils::tofNSigma(o2::track::PID::Pion, track), track.pt());
+          registry.fill(HIST("hTpcTofNSigmaPidKaonMc"), o2::aod::pidutils::tpcNSigma(o2::track::PID::Kaon, track), o2::aod::pidutils::tofNSigma(o2::track::PID::Kaon, track), track.pt());
+          registry.fill(HIST("hTpcTofNSigmaPidProtonMc"), o2::aod::pidutils::tpcNSigma(o2::track::PID::Proton, track), o2::aod::pidutils::tofNSigma(o2::track::PID::Proton, track), track.pt());
+        }
         bool isPhysicalPrimary = false;
         // DsToKKPi and DsToPiKK division
         if (isCorrectInvMassHypo && candidate.isSelDsToKKPi() >= selectionFlagDs) {
-          entryDsHadronPair(getDeltaPhi(track.phi(), candidate.phi()),
-                            track.eta() - candidate.eta(),
-                            candidate.pt() * chargeDs,
-                            track.pt() * track.sign(),
-                            poolBin,
-                            collision.numContrib());
-          entryDsHadronRecoInfo(HfHelper::invMassDsToKKPi(candidate), isDsSignal, isDecayChan);
-          entryDsHadronMlInfo(outputMl[0], outputMl[2]);
           if (track.has_mcParticle()) {
             auto mcParticle = track.template mcParticle_as<aod::McParticles>();
+            if (trkPIDspecies->at(0) == o2::track::PID::Kaon && mcParticle.pdgCode() != kKPlus) {
+              continue;
+            }
+            if (trkPIDspecies->at(0) == o2::track::PID::Pion && mcParticle.pdgCode() != kPiPlus) {
+              continue;
+            }
+            if (trkPIDspecies->at(0) == o2::track::PID::Proton && mcParticle.pdgCode() != kProton) {
+              continue;
+            }
+            entryDsHadronPair(getDeltaPhi(track.phi(), candidate.phi()),
+                              track.eta() - candidate.eta(),
+                              candidate.pt() * chargeDs,
+                              track.pt() * track.sign(),
+                              poolBin,
+                              collision.numContrib());
+            entryDsHadronRecoInfo(HfHelper::invMassDsToKKPi(candidate), isDsSignal, isDecayChan);
+            entryDsHadronMlInfo(outputMl[0], outputMl[2]);
             isPhysicalPrimary = mcParticle.isPhysicalPrimary();
             auto trackOrigin = RecoDecay::getCharmHadronOrigin(mcParticles, mcParticle, true);
             entryDsHadronGenInfo(isDsPrompt, isPhysicalPrimary, trackOrigin);
           } else {
-            entryDsHadronGenInfo(isDsPrompt, isPhysicalPrimary, 0);
             registry.fill(HIST("hFakeTracksMcRec"), track.pt());
           }
           // for secondary particle fraction estimation
@@ -627,21 +644,20 @@ struct HfCorrelatorDsHadrons {
           }
           entryTrackRecoInfo(track.dcaXY(), track.dcaZ(), track.tpcNClsCrossedRows());
         } else if (isCorrectInvMassHypo && candidate.isSelDsToPiKK() >= selectionFlagDs) {
-          entryDsHadronPair(getDeltaPhi(track.phi(), candidate.phi()),
-                            track.eta() - candidate.eta(),
-                            candidate.pt() * chargeDs,
-                            track.pt() * track.sign(),
-                            poolBin,
-                            collision.numContrib());
-          entryDsHadronRecoInfo(HfHelper::invMassDsToPiKK(candidate), isDsSignal, isDecayChan);
-          entryDsHadronMlInfo(outputMl[0], outputMl[2]);
           if (track.has_mcParticle()) {
             auto mcParticle = track.template mcParticle_as<aod::McParticles>();
+            entryDsHadronPair(getDeltaPhi(track.phi(), candidate.phi()),
+                              track.eta() - candidate.eta(),
+                              candidate.pt() * chargeDs,
+                              track.pt() * track.sign(),
+                              poolBin,
+                              collision.numContrib());
+            entryDsHadronRecoInfo(HfHelper::invMassDsToPiKK(candidate), isDsSignal, isDecayChan);
+            entryDsHadronMlInfo(outputMl[0], outputMl[2]);
             isPhysicalPrimary = mcParticle.isPhysicalPrimary();
             auto trackOrigin = RecoDecay::getCharmHadronOrigin(mcParticles, mcParticle, true);
             entryDsHadronGenInfo(isDsPrompt, isPhysicalPrimary, trackOrigin);
           } else {
-            entryDsHadronGenInfo(isDsPrompt, false, 0);
             registry.fill(HIST("hFakeTracksMcRec"), track.pt());
           }
           // for secondary particle fraction estimation

PWGHF/HFC/TableProducer/correlatorFlowCharmHadronsReduced.cxx

@@ -155,7 +155,7 @@ struct HfCorrelatorFlowCharmHadronsReduced {
       LOGP(fatal, "The size of bdtScore0PtMaxs and bdtScore1PtMins must be the one of binsPtTrig minus one!");
     }
 
-    if (doprocessSameEventCharmHadWCentMix || doprocessSameEventHadHadWCentMix || doprocessMixedEventCharmHadWCentMix || doprocessMixedEventHadHadWCentMix) {
+    if (doprocessSameEventCharmHadWCentMix || doprocessSameEventHadHadWCentMix || doprocessMixedEventCharmHadWCentMix || doprocessMixedEventHadHadWCentMix || doprocessSameEventCharmHadWCentMixBase || doprocessMixedEventCharmHadWCentMixBase) {
       poolBins = (centPoolBins->size() - 2) * (zPoolBins->size() - 2);
     } else {
       poolBins = (multPoolBins->size() - 2) * (zPoolBins->size() - 2);
@@ -434,8 +434,8 @@ struct HfCorrelatorFlowCharmHadronsReduced {
         continue;
       }
 
-      int poolBin = binPolicy.getBin({c2.posZ(), c2.multiplicity()});
-      int poolBinTrigCand = binPolicy.getBin({c1.posZ(), c1.multiplicity()});
+      int poolBin = binPolicy.getBin({c2.posZ(), c2.centrality()});
+      int poolBinTrigCand = binPolicy.getBin({c1.posZ(), c1.centrality()});
 
       if (poolBin != poolBinTrigCand) {
         LOGF(info, "Error, poolBins are different");
@@ -618,7 +618,7 @@ struct HfCorrelatorFlowCharmHadronsReduced {
       if (collision.centrality() < centralityMin || collision.centrality() > centralityMax) {
         continue;
       }
-      int poolBin = binPolicyPosZCent.getBin({collision.posZ(), collision.multiplicity()});
+      int poolBin = binPolicyPosZCent.getBin({collision.posZ(), collision.centrality()});
 
       auto thisCollId = collision.globalIndex();
       auto candsThisColl = candidates.sliceBy(trigCharmCandsPerCol, thisCollId);

PWGHF/HFC/Tasks/taskCorrelationDsHadrons.cxx

@@ -14,6 +14,7 @@
 /// \author Grazia Luparello <grazia.luparello@cern.ch>
 /// \author Samuele Cattaruzzi <samuele.cattaruzzi@cern.ch>
 
+#include "PWGHF/Core/CentralityEstimation.h"
 #include "PWGHF/Core/DecayChannels.h"
 #include "PWGHF/Core/HfHelper.h"
 #include "PWGHF/Core/SelectorCuts.h"
@@ -22,10 +23,12 @@
 #include "PWGHF/DataModel/CandidateSelectionTables.h"
 #include "PWGHF/DataModel/TrackIndexSkimmingTables.h"
 #include "PWGHF/HFC/DataModel/CorrelationTables.h"
+#include "PWGHF/HFC/Utils/utilsCorrelations.h"
 #include "PWGHF/Utils/utilsAnalysis.h"
 
 #include "Common/CCDB/EventSelectionParams.h"
 #include "Common/Core/RecoDecay.h"
+#include "Common/DataModel/Centrality.h"
 #include "Common/DataModel/EventSelection.h"
 #include "Common/DataModel/Multiplicity.h"
 #include "Common/DataModel/TrackSelectionTables.h"
@@ -65,6 +68,8 @@ using namespace o2::constants::physics;
 using namespace o2::constants::math;
 using namespace o2::framework;
 using namespace o2::framework::expressions;
+using namespace o2::hf_centrality;
+using namespace o2::analysis::hf_correlations;
 
 enum ResonantChannel : int8_t {
   PhiPi = 1,
@@ -92,7 +97,10 @@ struct HfTaskCorrelationDsHadrons {
   Configurable<bool> applyDeltaPhiCorrEff{"applyDeltaPhiCorrEff", false, "Flag to use higher dimension histograms with delta phi correction in the efficiency correction"};
   Configurable<bool> doLSpair{"doLSpair", false, "Flag to evaluate correlations for like-sign pairs"};
   Configurable<bool> doULSpair{"doULSpair", false, "Flag to evaluate correlations for unlike-sign pairs"};
+  Configurable<bool> pidTrkApplied{"pidTrkApplied", false, "Apply PID selection for associated tracks"};
+  Configurable<bool> forceTOF{"forceTOF", false, "force the TOF signal for the PID"};
   // Configurable<bool> doMcCollisionCheck{"doMcCollisionCheck", false, "Flag for applying the collision check and selection based on MC collision info"};
+  Configurable<int> centEstimator{"centEstimator", 2, "Centrality estimation (FT0A: 1, FT0C: 2, FT0M: 3, FV0A: 4)"};
   Configurable<int> selectionFlagDs{"selectionFlagDs", 7, "Selection Flag for Ds (avoid the case of flag = 0, no outputMlScore)"};
   Configurable<int> nTpcCrossedRaws{"nTpcCrossedRaws", 70, "Number of crossed TPC Rows"};
   // Configurable<int> eventGeneratorType{"eventGeneratorType", -1, "If positive, enable event selection using subGeneratorId information. The value indicates which events to keep (0 = MB, 4 = charm triggered, 5 = beauty triggered)"};
@@ -108,7 +116,13 @@ struct HfTaskCorrelationDsHadrons {
   Configurable<float> yCandMax{"yCandMax", 0.8, "max. cand. rapidity (used in eff. process only)"};
   Configurable<float> yCandGenMax{"yCandGenMax", 0.5, "max. gen. cand. rapidity (used in eff. process only)"};
   Configurable<float> ptDaughterMin{"ptDaughterMin", 0.1, "min. daughter pT (used in eff. process only)"};
+  Configurable<float> tofPIDThreshold{"tofPIDThreshold", 0.75, "minimum pT after which TOF PID is applicable"};
+  Configurable<float> centralityMin{"centralityMin", 0., "min. centrality"};
+  Configurable<float> centralityMax{"centralityMax", 100., "max. centrality"};
   Configurable<std::vector<int>> classMl{"classMl", {0, 1, 2}, "Indexes of ML scores to be stored. Three indexes max."};
+  Configurable<std::vector<int>> trkPIDspecies{"trkPIDspecies", std::vector<int>{o2::track::PID::Proton, o2::track::PID::Pion, o2::track::PID::Kaon}, "Trk sel: Particles species for PID, proton, pion, kaon"};
+  Configurable<std::vector<float>> pidTPCMax{"pidTPCMax", std::vector<float>{3., 0., 0.}, "maximum nSigma TPC"};
+  Configurable<std::vector<float>> pidTOFMax{"pidTOFMax", std::vector<float>{3., 0., 0.}, "maximum nSigma TOF"};
   Configurable<std::vector<double>> binsPtD{"binsPtD", std::vector<double>{o2::analysis::hf_cuts_ds_to_k_k_pi::vecBinsPt}, "pT bin limits for candidate mass plots and efficiency"};
   Configurable<std::vector<double>> binsPtHadron{"binsPtHadron", std::vector<double>{0.3, 2., 4., 8., 12., 50.}, "pT bin limits for assoc particle efficiency"};
   Configurable<std::vector<double>> mlOutputPromptMin{"mlOutputPromptMin", {0.5, 0.5, 0.5, 0.5}, "Machine learning scores for prompt"};
@@ -143,16 +157,17 @@ struct HfTaskCorrelationDsHadrons {
 
   Service<ccdb::BasicCCDBManager> ccdb;
 
-  using DsHadronPair = soa::Join<aod::DsHadronPair, aod::DsHadronRecoInfo>;
-  using DsHadronPairFull = soa::Join<aod::DsHadronPair, aod::DsHadronRecoInfo, aod::DsHadronGenInfo>;
-  using DsHadronPairWithMl = soa::Join<aod::DsHadronPair, aod::DsHadronRecoInfo, aod::DsHadronMlInfo, aod::TrackRecoInfo>;
-  using DsHadronPairFullWithMl = soa::Join<aod::DsHadronPair, aod::DsHadronRecoInfo, aod::DsHadronGenInfo, aod::DsHadronMlInfo, aod::TrackRecoInfo>;
+  using DsHadronPair = soa::Filtered<soa::Join<aod::DsHadronPair, aod::DsHadronRecoInfo>>;
+  using DsHadronPairFull = soa::Filtered<soa::Join<aod::DsHadronPair, aod::DsHadronRecoInfo, aod::DsHadronGenInfo>>;
+  using DsHadronPairWithMl = soa::Filtered<soa::Join<aod::DsHadronPair, aod::DsHadronRecoInfo, aod::DsHadronMlInfo, aod::TrackRecoInfo>>;
+  using DsHadronPairFullWithMl = soa::Filtered<soa::Join<aod::DsHadronPair, aod::DsHadronRecoInfo, aod::DsHadronGenInfo, aod::DsHadronMlInfo, aod::TrackRecoInfo>>;
   using CandDsMcReco = soa::Filtered<soa::Join<aod::HfCand3Prong, aod::HfSelDsToKKPi, aod::HfMlDsToKKPi, aod::HfCand3ProngMcRec>>; // flagDsFilter applied
   using CandDsMcGen = soa::Join<aod::McParticles, aod::HfCand3ProngMcGen>;                                                         // flagDsFilter applied
   using TracksWithMc = soa::Filtered<soa::Join<aod::TracksWDca, aod::TrackSelection, aod::TracksExtra, o2::aod::McTrackLabels>>;   // trackFilter applied
 
   Filter flagDsFilter = ((o2::aod::hf_track_index::hfflag & static_cast<uint8_t>(1 << aod::hf_cand_3prong::DecayType::DsToKKPi)) != static_cast<uint8_t>(0)) && (aod::hf_sel_candidate_ds::isSelDsToKKPi >= selectionFlagDs || aod::hf_sel_candidate_ds::isSelDsToPiKK >= selectionFlagDs);
   Filter trackFilter = (nabs(aod::track::eta) < etaTrackMax) && (aod::track::pt > ptTrackMin) && (aod::track::pt < ptTrackMax) && (nabs(aod::track::dcaXY) < dcaXYTrackMax) && (nabs(aod::track::dcaZ) < dcaZTrackMax);
+  Filter trackPairFilter = (nabs(aod::hf_correlation_ds_hadron::signedPtHadron) > ptTrackMin) && (nabs(aod::hf_correlation_ds_hadron::signedPtHadron) < ptTrackMax);
 
   Preslice<CandDsMcReco> perCollisionCand = o2::aod::hf_cand::collisionId;
   Preslice<CandDsMcGen> perCollisionCandMc = o2::aod::mcparticle::mcCollisionId;
@@ -431,6 +446,27 @@ struct HfTaskCorrelationDsHadrons {
     return weight;
   }
 
+  /// Check event selections for collision and fill the collision table
+  /// \param collision is the collision
+  template <typename Coll>
+  int getCentrality(Coll const& collision)
+  {
+    int cent{-1};
+    if (centEstimator == CentralityEstimator::FT0A) {
+      cent = collision.centFT0A();
+    } else if (centEstimator == CentralityEstimator::FT0C) {
+      cent = collision.centFT0C();
+    } else if (centEstimator == CentralityEstimator::FT0M) {
+      cent = collision.centFT0M();
+    } else if (centEstimator == CentralityEstimator::FV0A) {
+      cent = collision.centFV0A();
+    } else {
+      LOG(fatal) << "Centrality estimator not recognized for collision selection";
+      std::abort();
+    }
+    return cent;
+  }
+
   void processData(DsHadronPairWithMl const& pairEntries,
                    aod::DsCandRecoInfo const& candidates)
   {
@@ -1061,7 +1097,7 @@ struct HfTaskCorrelationDsHadrons {
   PROCESS_SWITCH(HfTaskCorrelationDsHadrons, processMcCandEfficiencyWoColl, "Process MC for calculating candidate reconstruction efficiency", false);
 
   /// Ds-Hadron correlation - for calculating associated particle tracking efficiency using MC reco-level analysis
-  void processMcTrackEfficiency(soa::Join<aod::Collisions, aod::FT0Mults, aod::EvSels, aod::McCollisionLabels> const& collisions,
+  void processMcTrackEfficiency(soa::Join<aod::Collisions, aod::FT0Mults, aod::EvSels, aod::CentFT0Ms, aod::CentFT0As, aod::CentFT0Cs, aod::CentFV0As, aod::McCollisionLabels> const& collisions,
                                 soa::Join<aod::McCollisions, aod::MultsExtraMC> const& mcCollisions,
                                 aod::McParticles const& mcParticles,
                                 TracksWithMc const& tracksData)
@@ -1093,6 +1129,10 @@ struct HfTaskCorrelationDsHadrons {
         if (selNoSameBunchPileUpColl && !(collision.selection_bit(o2::aod::evsel::kNoSameBunchPileup))) {
           continue;
         }
+        int centrality = getCentrality(collision);
+        if (centrality < centralityMin || centrality > centralityMax) {
+          continue;
+        }
         if (!collision.has_mcCollision()) {
           registry.fill(HIST("hFakeCollision"), 0.);
           continue;
@@ -1151,6 +1191,11 @@ struct HfTaskCorrelationDsHadrons {
             continue;
           }
           if (track.has_mcParticle()) {
+            if (pidTrkApplied) {
+              if (!passPIDSelection(track, trkPIDspecies, pidTPCMax, pidTOFMax, tofPIDThreshold, forceTOF)) {
+                continue;
+              }
+            }
             auto mcParticle = track.template mcParticle_as<aod::McParticles>();
             if (mcParticle.isPhysicalPrimary()) {
               registry.fill(HIST("hPtParticleAssocMcRec"), track.pt());