[PWGHF] Added centrality to D+ tree tables

PWGHF/TableProducer/treeCreatorDplusToPiKPi.cxx

@@ -16,17 +16,21 @@
 ///
 /// \author Alexandre Bigot <alexandre.bigot@cern.ch>, IPHC Strasbourg
 
+#include <vector>
+
 #include "CommonConstants/PhysicsConstants.h"
 #include "Framework/AnalysisTask.h"
 #include "Framework/runDataProcessing.h"
 
 #include "PWGHF/Core/HfHelper.h"
+#include "PWGHF/Core/CentralityEstimation.h"
 #include "PWGHF/DataModel/CandidateReconstructionTables.h"
 #include "PWGHF/DataModel/CandidateSelectionTables.h"
 
 using namespace o2;
 using namespace o2::framework;
 using namespace o2::framework::expressions;
+using namespace o2::hf_centrality;
 
 namespace o2::aod
 {
@@ -50,6 +54,7 @@ DECLARE_SOA_COLUMN(Y, y, float);
 DECLARE_SOA_COLUMN(Eta, eta, float);                                               //! Pseudorapidity of candidate
 DECLARE_SOA_COLUMN(Phi, phi, float);                                               //! Azimuth angle of candidate
 DECLARE_SOA_COLUMN(E, e, float);                                                   //! Energy of candidate (GeV)
+DECLARE_SOA_COLUMN(Centrality, centrality, float);                                 //! Collision centrality
 DECLARE_SOA_COLUMN(NSigTpcPi0, nSigTpcPi0, float);                                 //! TPC Nsigma separation for prong0 with pion mass hypothesis
 DECLARE_SOA_COLUMN(NSigTpcKa0, nSigTpcKa0, float);                                 //! TPC Nsigma separation for prong0 with kaon mass hypothesis
 DECLARE_SOA_COLUMN(NSigTofPi0, nSigTofPi0, float);                                 //! TOF Nsigma separation for prong0 with pion mass hypothesis
@@ -129,6 +134,7 @@ DECLARE_SOA_TABLE(HfCandDpLites, "AOD", "HFCANDDPLITE",
                   full::Eta,
                   full::Phi,
                   full::Y,
+                  full::Centrality,
                   hf_cand_3prong::FlagMcMatchRec,
                   hf_cand_3prong::OriginMcRec,
                   hf_cand_3prong::FlagMcDecayChanRec)
@@ -210,6 +216,7 @@ DECLARE_SOA_TABLE(HfCandDpFulls, "AOD", "HFCANDDPFULL",
                   full::Phi,
                   full::Y,
                   full::E,
+                  full::Centrality,
                   hf_cand_3prong::FlagMcMatchRec,
                   hf_cand_3prong::OriginMcRec,
                   hf_cand_3prong::FlagMcDecayChanRec);
@@ -249,7 +256,8 @@ struct HfTreeCreatorDplusToPiKPi {
   Configurable<bool> fillOnlyBackground{"fillOnlyBackground", false, "Flag to fill derived tables with background for ML trainings"};
   Configurable<float> downSampleBkgFactor{"downSampleBkgFactor", 1., "Fraction of background candidates to keep for ML trainings"};
   Configurable<float> ptMaxForDownSample{"ptMaxForDownSample", 10., "Maximum pt for the application of the downsampling factor"};
-  Configurable<std::vector<int>> classMl{"classMlindexes", {0, 2}, "Indexes of ML bkg and non-prompt scores."};
+  Configurable<std::vector<int>> classMlIndexes{"classMlIndexes", {0, 2}, "Indexes of ML bkg and non-prompt scores."};
+  Configurable<int> centEstimator{"centEstimator", 0, "Centrality estimation (None: 0, FT0C: 2, FT0M: 3)"};
 
   HfHelper hfHelper;
 
@@ -258,6 +266,8 @@ struct HfTreeCreatorDplusToPiKPi {
   using SelectedCandidatesMcWithMl = soa::Filtered<soa::Join<aod::HfCand3Prong, aod::HfCand3ProngMcRec, aod::HfSelDplusToPiKPi, aod::HfMlDplusToPiKPi>>;
   using TracksWPid = soa::Join<aod::Tracks, aod::TracksPidPi, aod::PidTpcTofFullPi, aod::TracksPidKa, aod::PidTpcTofFullKa>;
 
+  using CollisionsCent = soa::Join<aod::Collisions, aod::CentFT0Cs, aod::CentFT0Ms>;
+
   Filter filterSelectCandidates = aod::hf_sel_candidate_dplus::isSelDplusToPiKPi >= selectionFlagDplus;
   Filter filterMcGenMatching = nabs(o2::aod::hf_cand_3prong::flagMcMatchGen) == static_cast<int8_t>(BIT(aod::hf_cand_3prong::DecayType::DplusToPiKPi));
 
@@ -282,7 +292,7 @@ struct HfTreeCreatorDplusToPiKPi {
       runNumber);
   }
 
-  template <bool doMc = false, bool doMl = false, typename T>
+  template <typename Coll, bool doMc = false, bool doMl = false, typename T>
   void fillCandidateTable(const T& candidate)
   {
     int8_t flagMc = 0;
@@ -296,8 +306,8 @@ struct HfTreeCreatorDplusToPiKPi {
 
     std::vector<float> outputMl = {-999., -999.};
     if constexpr (doMl) {
-      for (unsigned int iclass = 0; iclass < classMl->size(); iclass++) {
-        outputMl[iclass] = candidate.mlProbDplusToPiKPi()[classMl->at(iclass)];
+      for (unsigned int iclass = 0; iclass < classMlIndexes->size(); iclass++) {
+        outputMl[iclass] = candidate.mlProbDplusToPiKPi()[classMlIndexes->at(iclass)];
       }
       rowCandidateMl(
         outputMl[0],
@@ -308,6 +318,12 @@ struct HfTreeCreatorDplusToPiKPi {
     auto prong1 = candidate.template prong1_as<TracksWPid>();
     auto prong2 = candidate.template prong2_as<TracksWPid>();
 
+    float cent{-1.};
+    auto coll = candidate.template collision_as<Coll>();
+    if (std::is_same_v<Coll, CollisionsCent> && centEstimator != CentralityEstimator::None) {
+      cent = getCentralityColl(coll, centEstimator);
+    }
+
     if (fillCandidateLiteTable) {
       rowCandidateLite(
         candidate.chi2PCA(),
@@ -351,13 +367,14 @@ struct HfTreeCreatorDplusToPiKPi {
         candidate.eta(),
         candidate.phi(),
         hfHelper.yDplus(candidate),
+        cent,
         flagMc,
         originMc,
         channelMc);
     } else {
       rowCandidateFull(
-        candidate.collision().bcId(),
-        candidate.collision().numContrib(),
+        coll.bcId(),
+        coll.numContrib(),
         candidate.posX(),
         candidate.posY(),
         candidate.posZ(),
@@ -432,6 +449,7 @@ struct HfTreeCreatorDplusToPiKPi {
         candidate.phi(),
         hfHelper.yDplus(candidate),
         hfHelper.eDplus(candidate),
+        cent,
         flagMc,
         originMc,
         channelMc);
@@ -461,7 +479,7 @@ struct HfTreeCreatorDplusToPiKPi {
           continue;
         }
       }
-      fillCandidateTable(candidate);
+      fillCandidateTable<aod::Collisions>(candidate);
     }
   }
 
@@ -488,7 +506,7 @@ struct HfTreeCreatorDplusToPiKPi {
         rowCandidateFull.reserve(reconstructedCandSig.size());
       }
       for (const auto& candidate : reconstructedCandSig) {
-        fillCandidateTable<true>(candidate);
+        fillCandidateTable<aod::Collisions, true>(candidate);
       }
     } else if (fillOnlySignalMl) {
       rowCandidateMl.reserve(reconstructedCandSigMl.size());
@@ -499,12 +517,12 @@ struct HfTreeCreatorDplusToPiKPi {
       }
       for (const auto& candidate : reconstructedCandSigMl) {
         if (downSampleBkgFactor < 1.) {
-          float pseudoRndm = candidate.ptProng0() * 1000. - (int64_t)(candidate.ptProng0() * 1000);
+          float pseudoRndm = candidate.ptProng0() * 1000. - static_cast<int64_t>(candidate.ptProng0() * 1000);
           if (candidate.pt() < ptMaxForDownSample && pseudoRndm >= downSampleBkgFactor) {
             continue;
           }
         }
-        fillCandidateTable<true, true>(candidate);
+        fillCandidateTable<aod::Collisions, true, true>(candidate);
       }
     } else if (fillOnlyBackground) {
       if (fillCandidateLiteTable) {
@@ -519,7 +537,7 @@ struct HfTreeCreatorDplusToPiKPi {
             continue;
           }
         }
-        fillCandidateTable<true>(candidate);
+        fillCandidateTable<aod::Collisions, true>(candidate);
       }
     } else {
       if (fillCandidateLiteTable) {
@@ -528,7 +546,7 @@ struct HfTreeCreatorDplusToPiKPi {
         rowCandidateFull.reserve(candidates.size());
       }
       for (const auto& candidate : candidates) {
-        fillCandidateTable<true>(candidate);
+        fillCandidateTable<aod::Collisions, true>(candidate);
       }
     }
 
@@ -547,6 +565,116 @@ struct HfTreeCreatorDplusToPiKPi {
   }
 
   PROCESS_SWITCH(HfTreeCreatorDplusToPiKPi, processMc, "Process MC", false);
+
+  void processDataWCent(CollisionsCent const& collisions,
+                        soa::Filtered<soa::Join<aod::HfCand3Prong, aod::HfSelDplusToPiKPi>> const& candidates,
+                        TracksWPid const&)
+  {
+    // Filling event properties
+    rowCandidateFullEvents.reserve(collisions.size());
+    for (const auto& collision : collisions) {
+      fillEvent(collision, 0, 1);
+    }
+
+    // Filling candidate properties
+    if (fillCandidateLiteTable) {
+      rowCandidateLite.reserve(candidates.size());
+    } else {
+      rowCandidateFull.reserve(candidates.size());
+    }
+    for (const auto& candidate : candidates) {
+      if (downSampleBkgFactor < 1.) {
+        float pseudoRndm = candidate.ptProng0() * 1000. - static_cast<int64_t>(candidate.ptProng0() * 1000);
+        if (candidate.pt() < ptMaxForDownSample && pseudoRndm >= downSampleBkgFactor) {
+          continue;
+        }
+      }
+      fillCandidateTable<CollisionsCent>(candidate);
+    }
+  }
+
+  PROCESS_SWITCH(HfTreeCreatorDplusToPiKPi, processDataWCent, "Process data with cent", false);
+
+  void processMcWCent(CollisionsCent const& collisions,
+                      aod::McCollisions const&,
+                      SelectedCandidatesMc const& candidates,
+                      MatchedGenCandidatesMc const& particles,
+                      SelectedCandidatesMcWithMl const&,
+                      TracksWPid const&)
+  {
+    // Filling event properties
+    rowCandidateFullEvents.reserve(collisions.size());
+    for (const auto& collision : collisions) {
+      fillEvent(collision, 0, 1);
+    }
+
+    // Filling candidate properties
+    if (fillOnlySignal) {
+      if (fillCandidateLiteTable) {
+        rowCandidateLite.reserve(reconstructedCandSig.size());
+      } else {
+        rowCandidateFull.reserve(reconstructedCandSig.size());
+      }
+      for (const auto& candidate : reconstructedCandSig) {
+        fillCandidateTable<CollisionsCent, true>(candidate);
+      }
+    } else if (fillOnlySignalMl) {
+      rowCandidateMl.reserve(reconstructedCandSigMl.size());
+      if (fillCandidateLiteTable) {
+        rowCandidateLite.reserve(reconstructedCandSigMl.size());
+      } else {
+        rowCandidateFull.reserve(reconstructedCandSigMl.size());
+      }
+      for (const auto& candidate : reconstructedCandSigMl) {
+        if (downSampleBkgFactor < 1.) {
+          float pseudoRndm = candidate.ptProng0() * 1000. - static_cast<int64_t>(candidate.ptProng0() * 1000);
+          if (candidate.pt() < ptMaxForDownSample && pseudoRndm >= downSampleBkgFactor) {
+            continue;
+          }
+        }
+        fillCandidateTable<CollisionsCent, true, true>(candidate);
+      }
+    } else if (fillOnlyBackground) {
+      if (fillCandidateLiteTable) {
+        rowCandidateLite.reserve(reconstructedCandBkg.size());
+      } else {
+        rowCandidateFull.reserve(reconstructedCandBkg.size());
+      }
+      for (const auto& candidate : reconstructedCandBkg) {
+        if (downSampleBkgFactor < 1.) {
+          float pseudoRndm = candidate.ptProng0() * 1000. - static_cast<int64_t>(candidate.ptProng0() * 1000);
+          if (candidate.pt() < ptMaxForDownSample && pseudoRndm >= downSampleBkgFactor) {
+            continue;
+          }
+        }
+        fillCandidateTable<CollisionsCent, true>(candidate);
+      }
+    } else {
+      if (fillCandidateLiteTable) {
+        rowCandidateLite.reserve(candidates.size());
+      } else {
+        rowCandidateFull.reserve(candidates.size());
+      }
+      for (const auto& candidate : candidates) {
+        fillCandidateTable<CollisionsCent, true>(candidate);
+      }
+    }
+
+    // Filling particle properties
+    rowCandidateFullParticles.reserve(particles.size());
+    for (const auto& particle : particles) {
+      rowCandidateFullParticles(
+        particle.mcCollision().bcId(),
+        particle.pt(),
+        particle.eta(),
+        particle.phi(),
+        RecoDecay::y(particle.pVector(), o2::constants::physics::MassDPlus),
+        particle.flagMcMatchGen(),
+        particle.originMcGen());
+    }
+  }
+
+  PROCESS_SWITCH(HfTreeCreatorDplusToPiKPi, processMcWCent, "Process MC with cent", false);
 };
 
 WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)