implement table producing for p-Dplus femto

Yunfan-Liu · Yunfan-Liu · commit c845ad552bab · 2025-07-18T10:21:30.000+08:00
diff --git a/PWGHF/HFC/TableProducer/femtoDreamProducer.cxx b/PWGHF/HFC/TableProducer/femtoDreamProducer.cxx
@@ -21,6 +21,7 @@
 #include "PWGHF/Core/CentralityEstimation.h"
 #include "PWGHF/Core/DecayChannels.h"
 #include "PWGHF/Core/HfHelper.h"
+#include "PWGHF/Core/HfMlResponseDplusToPiKPi.h"
 #include "PWGHF/Core/HfMlResponseLcToPKPi.h"
 #include "PWGHF/Core/SelectorCuts.h"
 #include "PWGHF/DataModel/CandidateReconstructionTables.h"
@@ -87,6 +88,11 @@ enum MlMode : uint8_t {
   FillMlFromNewBDT
 };
 
+// decay channels
+enum DecayChannel { DplusToPiKPi = 0,
+                    LcToPKPi
+};
+
 struct HfFemtoDreamProducer {
 
   Produces<aod::FDCollisions> outputCollision;
@@ -119,9 +125,9 @@ struct HfFemtoDreamProducer {
   Configurable<bool> isDebug{"isDebug", true, "Enable Debug tables"};
   Configurable<bool> isRun3{"isRun3", true, "Running on Run3 or pilot"};
 
-  /// Lc table
-  Configurable<int> selectionFlagLc{"selectionFlagLc", 1, "Selection Flag for Lc"};
-  Configurable<bool> useCent{"useCent", false, "Enable centrality for lc"};
+  /// Charm hadron table
+  Configurable<int> selectionFlagHadron{"selectionFlagHadron", 1, "Selection Flag for Charm Hadron"};
+  Configurable<bool> useCent{"useCent", false, "Enable centrality for Charm Hadron"};
 
   Configurable<int> trkPDGCode{"trkPDGCode", 2212, "PDG code of the selected track for Monte Carlo truth"};
   Configurable<std::vector<float>> trkCharge{FemtoDreamTrackSelection::getSelectionName(femtoDreamTrackSelection::kSign, "trk"), std::vector<float>{-1, 1}, FemtoDreamTrackSelection::getSelectionHelper(femtoDreamTrackSelection::kSign, "Track selection: ")};
@@ -141,7 +147,7 @@ struct HfFemtoDreamProducer {
   Configurable<std::vector<float>> trkITSnclsIbMin{FemtoDreamTrackSelection::getSelectionName(femtoDreamTrackSelection::kITSnClsIbMin, "trk"), std::vector<float>{-1.f, 1.f}, FemtoDreamTrackSelection::getSelectionHelper(femtoDreamTrackSelection::kITSnClsIbMin, "Track selection: ")};
   Configurable<std::vector<float>> trkITSnclsMin{FemtoDreamTrackSelection::getSelectionName(femtoDreamTrackSelection::kITSnClsMin, "trk"), std::vector<float>{-1.f, 2.f, 4.f}, FemtoDreamTrackSelection::getSelectionHelper(femtoDreamTrackSelection::kITSnClsMin, "Track selection: ")};
   // ML inference
-  Configurable<int> applyMlMode{"applyMlMode", 1, "None: 0, BDT model from Lc selector: 1, New BDT model on Top of Lc selector: 2"};
+  Configurable<int> applyMlMode{"applyMlMode", 1, "None: 0, BDT model from candidate selector: 1, New BDT model on Top of candidate selector: 2"};
   Configurable<std::vector<double>> binsPtMl{"binsPtMl", std::vector<double>{hf_cuts_ml::vecBinsPt}, "pT bin limits for ML application"};
   Configurable<std::vector<int>> cutDirMl{"cutDirMl", std::vector<int>{hf_cuts_ml::vecCutDir}, "Whether to reject score values greater or smaller than the threshold"};
   Configurable<LabeledArray<double>> cutsMl{"cutsMl", {hf_cuts_ml::Cuts[0], hf_cuts_ml::NBinsPt, hf_cuts_ml::NCutScores, hf_cuts_ml::labelsPt, hf_cuts_ml::labelsCutScore}, "ML selections per pT bin"};
@@ -152,6 +158,8 @@ struct HfFemtoDreamProducer {
 
   HfHelper hfHelper;
   o2::analysis::HfMlResponseLcToPKPi<float> hfMlResponse;
+  o2::analysis::HfMlResponseDplusToPiKPi<float> hfMlResponseDplus;
+  std::vector<float> outputMlDplus = {};
   std::vector<float> outputMlPKPi = {};
   std::vector<float> outputMlPiKP = {};
   o2::ccdb::CcdbApi ccdbApi;
@@ -162,6 +170,8 @@ struct HfFemtoDreamProducer {
 
   float magField;
   int runNumber;
+  using CandidateDplus = soa::Join<aod::HfCand3ProngWPidPiKa, aod::HfSelDplusToPiKPi>;
+  using CandidateDplusMc = soa::Join<aod::HfCand3ProngWPidPiKa, aod::HfSelDplusToPiKPi, aod::HfCand3ProngMcRec>;
   using CandidateLc = soa::Join<aod::HfCand3ProngWPidPiKaPr, aod::HfSelLc>;
   using CandidateLcMc = soa::Join<aod::HfCand3ProngWPidPiKaPr, aod::HfSelLc, aod::HfCand3ProngMcRec>;
 
@@ -176,14 +186,16 @@ struct HfFemtoDreamProducer {
 
   using GeneratedMc = soa::Filtered<soa::Join<aod::McParticles, aod::HfCand3ProngMcGen>>;
 
-  Filter filterSelectCandidateLc = (aod::hf_sel_candidate_lc::isSelLcToPKPi >= selectionFlagLc || aod::hf_sel_candidate_lc::isSelLcToPiKP >= selectionFlagLc);
+  Filter filterSelectCandidateDplus = (aod::hf_sel_candidate_dplus::isSelDplusToPiKPi >= selectionFlagHadron  || aod::hf_sel_candidate_dplus::isSelDplusToPiKPi >= selectionFlagHadron);
+  Filter filterSelectCandidateLc = (aod::hf_sel_candidate_lc::isSelLcToPKPi >= selectionFlagHadron || aod::hf_sel_candidate_lc::isSelLcToPiKP >= selectionFlagHadron);
 
   HistogramRegistry qaRegistry{"QAHistos", {}, OutputObjHandlingPolicy::AnalysisObject};
   HistogramRegistry trackRegistry{"Tracks", {}, OutputObjHandlingPolicy::AnalysisObject};
 
   void init(InitContext&)
   {
-    std::array<bool, 5> processes = {doprocessDataCharmHad, doprocessMcCharmHad, doprocessDataCharmHadWithML, doprocessMcCharmHadWithML, doprocessMcCharmHadGen};
+    std::array<bool, 10> processes = {doprocessDataDplusToPiKPi, doprocessMcDplusToPiKPi, doprocessDataDplusToPiKPiWithML, doprocessMcDplusToPiKPiWithML, doprocessMcDplusToPiKPiGen,
+                                      doprocessDataLcToPKPi, doprocessMcLcToPKPi, doprocessDataLcToPKPiWithML, doprocessMcLcToPKPiWithML, doprocessMcLcToPKPiGen};
     if (std::accumulate(processes.begin(), processes.end(), 0) != 1) {
       LOGP(fatal, "One and only one process function must be enabled at a time.");
     }
@@ -404,7 +416,7 @@ struct HfFemtoDreamProducer {
     return fIsTrackFilled;
   }
 
-  template <bool isMc, bool useCharmMl, typename TrackType, typename CollisionType, typename CandType>
+  template <DecayChannel channel, bool isMc, bool useCharmMl, typename TrackType, typename CollisionType, typename CandType>
   void fillCharmHadronTable(CollisionType const& col, TrackType const& tracks, CandType const& candidates)
   {
     const auto vtxZ = col.posZ();
@@ -451,44 +463,15 @@ struct HfFemtoDreamProducer {
     bool isTrackFilled = false;
     bool isSelectedMlLcToPKPi = true;
     bool isSelectedMlLcToPiKP = true;
+    bool isSelectedMlDplusToPiKPi = true;
     for (const auto& candidate : candidates) {
+      outputMlDplus = {-1.0f, -1.0f, -1.0f};
       outputMlPKPi = {-1.0f, -1.0f, -1.0f};
       outputMlPiKP = {-1.0f, -1.0f, -1.0f};
       auto trackPos1 = candidate.template prong0_as<TrackType>(); // positive daughter (negative for the antiparticles)
       auto trackNeg = candidate.template prong1_as<TrackType>();  // negative daughter (positive for the antiparticles)
       auto trackPos2 = candidate.template prong2_as<TrackType>(); // positive daughter (negative for the antiparticles)
-
-      if constexpr (useCharmMl) {
-        /// fill with ML information
-        /// BDT index 0: bkg score; BDT index 1: prompt score; BDT index 2: non-prompt score
-        if (applyMlMode == FillMlFromSelector) {
-          if (candidate.mlProbLcToPKPi().size() > 0) {
-            outputMlPKPi.at(0) = candidate.mlProbLcToPKPi()[0]; /// bkg score
-            outputMlPKPi.at(1) = candidate.mlProbLcToPKPi()[1]; /// prompt score
-            outputMlPKPi.at(2) = candidate.mlProbLcToPKPi()[2]; /// non-prompt score
-          }
-          if (candidate.mlProbLcToPiKP().size() > 0) {
-            outputMlPiKP.at(0) = candidate.mlProbLcToPiKP()[0]; /// bkg score
-            outputMlPiKP.at(1) = candidate.mlProbLcToPiKP()[1]; /// prompt score
-            outputMlPiKP.at(2) = candidate.mlProbLcToPiKP()[2]; /// non-prompt score
-          }
-        } else if (applyMlMode == FillMlFromNewBDT) {
-          isSelectedMlLcToPKPi = false;
-          isSelectedMlLcToPiKP = false;
-          if (candidate.mlProbLcToPKPi().size() > 0) {
-            std::vector<float> inputFeaturesLcToPKPi = hfMlResponse.getInputFeatures(candidate, true);
-            isSelectedMlLcToPKPi = hfMlResponse.isSelectedMl(inputFeaturesLcToPKPi, candidate.pt(), outputMlPKPi);
-          }
-          if (candidate.mlProbLcToPiKP().size() > 0) {
-            std::vector<float> inputFeaturesLcToPiKP = hfMlResponse.getInputFeatures(candidate, false);
-            isSelectedMlLcToPiKP = hfMlResponse.isSelectedMl(inputFeaturesLcToPiKP, candidate.pt(), outputMlPKPi);
-          }
-          if (!isSelectedMlLcToPKPi && !isSelectedMlLcToPiKP)
-            continue;
-        } else {
-          LOGF(fatal, "Please check your Ml configuration!!");
-        }
-      }
+      
       auto bc = col.template bc_as<aod::BCsWithTimestamps>();
       int64_t timeStamp = bc.timestamp();
       auto fillTable = [&](int CandFlag,
@@ -526,8 +509,65 @@ struct HfFemtoDreamProducer {
           }
       } };
 
-      fillTable(0, candidate.isSelLcToPKPi(), outputMlPKPi.at(0), outputMlPKPi.at(1), outputMlPKPi.at(2));
-      fillTable(1, candidate.isSelLcToPiKP(), outputMlPiKP.at(0), outputMlPiKP.at(1), outputMlPiKP.at(2));
+      if constexpr(channel == DecayChannel::DplusToPiKPi) {
+        if constexpr (useCharmMl) {
+          /// fill with ML information
+          /// BDT index 0: bkg score; BDT index 1: prompt score; BDT index 2: non-prompt score  
+          if (applyMlMode == FillMlFromSelector) {
+            if (candidate.mlProbDplusToPiKPi().size() > 0) {
+              outputMlDplus.at(0) = candidate.mlProbDplusToPiKPi()[0]; /// bkg score
+              outputMlDplus.at(1) = candidate.mlProbDplusToPiKPi()[1]; /// prompt score
+              outputMlDplus.at(2) = candidate.mlProbDplusToPiKPi()[2]; /// non-prompt score
+              }
+          } else if (applyMlMode == FillMlFromNewBDT) {
+              isSelectedMlDplusToPiKPi = false;
+              if (candidate.mlProbDplusToPiKPi().size() > 0) {
+                std::vector<float> inputFeaturesDplusToPiKPi = hfMlResponseDplus.getInputFeatures(candidate);
+                isSelectedMlDplusToPiKPi = hfMlResponseDplus.isSelectedMl(inputFeaturesDplusToPiKPi, candidate.pt(), outputMlDplus);
+              }
+              if (!isSelectedMlDplusToPiKPi)
+                continue;
+            } else {
+              LOGF(fatal, "Please check your Ml configuration!!");
+            }
+          }
+        fillTable(0, candidate.isSelDplusToPiKPi(), outputMlDplus.at(0), outputMlDplus.at(1), outputMlDplus.at(2));
+        
+      } else if constexpr(channel == DecayChannel::LcToPKPi) {
+        if constexpr (useCharmMl) {
+          /// fill with ML information
+          /// BDT index 0: bkg score; BDT index 1: prompt score; BDT index 2: non-prompt score  
+          if (applyMlMode == FillMlFromSelector) {
+            if (candidate.mlProbLcToPKPi().size() > 0) {
+              outputMlPKPi.at(0) = candidate.mlProbLcToPKPi()[0]; /// bkg score
+              outputMlPKPi.at(1) = candidate.mlProbLcToPKPi()[1]; /// prompt score
+              outputMlPKPi.at(2) = candidate.mlProbLcToPKPi()[2]; /// non-prompt score
+            }
+            if (candidate.mlProbLcToPiKP().size() > 0) {
+              outputMlPiKP.at(0) = candidate.mlProbLcToPiKP()[0]; /// bkg score
+              outputMlPiKP.at(1) = candidate.mlProbLcToPiKP()[1]; /// prompt score
+              outputMlPiKP.at(2) = candidate.mlProbLcToPiKP()[2]; /// non-prompt score
+            }
+          } else if (applyMlMode == FillMlFromNewBDT) {
+            isSelectedMlLcToPKPi = false;
+            isSelectedMlLcToPiKP = false;
+            if (candidate.mlProbLcToPKPi().size() > 0) {
+              std::vector<float> inputFeaturesLcToPKPi = hfMlResponse.getInputFeatures(candidate, true);
+              isSelectedMlLcToPKPi = hfMlResponse.isSelectedMl(inputFeaturesLcToPKPi, candidate.pt(), outputMlPKPi);
+            }
+            if (candidate.mlProbLcToPiKP().size() > 0) {
+              std::vector<float> inputFeaturesLcToPiKP = hfMlResponse.getInputFeatures(candidate, false);
+              isSelectedMlLcToPiKP = hfMlResponse.isSelectedMl(inputFeaturesLcToPiKP, candidate.pt(), outputMlPKPi);
+            }
+            if (!isSelectedMlLcToPKPi && !isSelectedMlLcToPiKP)
+              continue;
+          } else {
+            LOGF(fatal, "Please check your Ml configuration!!");
+          }
+        }
+        fillTable(0, candidate.isSelLcToPKPi(), outputMlPKPi.at(0), outputMlPKPi.at(1), outputMlPKPi.at(2));
+        fillTable(1, candidate.isSelLcToPiKP(), outputMlPiKP.at(0), outputMlPiKP.at(1), outputMlPiKP.at(2));
+      }
     }
 
     if (!isTrackFilled) {
@@ -586,20 +626,84 @@ struct HfFemtoDreamProducer {
     }
   }
 
-  void processDataCharmHad(FemtoFullCollision const& col,
+  /// DplusToPiKPi
+  void processDataDplusToPiKPi(FemtoFullCollision const& col,
+                           aod::BCsWithTimestamps const&,
+                           FemtoHFTracks const& tracks,
+                           soa::Filtered<CandidateDplus> const& candidates)
+  {
+    // get magnetic field for run
+    getMagneticFieldTesla(col.bc_as<aod::BCsWithTimestamps>());
+
+    fillCharmHadronTable<DecayChannel::DplusToPiKPi, false, false>(col, tracks, candidates);
+  }
+  PROCESS_SWITCH(HfFemtoDreamProducer, processDataDplusToPiKPi,
+                 "Provide experimental data for DplusToPiKPi femto", false);
+
+  void processDataDplusToPiKPiWithML(FemtoFullCollision const& col,
+                                 aod::BCsWithTimestamps const&,
+                                 FemtoHFTracks const& tracks,
+                                 soa::Filtered<soa::Join<CandidateDplus,
+                                                         aod::HfMlDplusToPiKPi>> const& candidates)
+  {
+
+    // get magnetic field for run
+    getMagneticFieldTesla(col.bc_as<aod::BCsWithTimestamps>());
+
+    fillCharmHadronTable<DecayChannel::DplusToPiKPi, false, true>(col, tracks, candidates);
+  }
+  PROCESS_SWITCH(HfFemtoDreamProducer, processDataDplusToPiKPiWithML,
+                 "Provide experimental data for DplusToPiKPi with ml", false);
+
+  void processMcDplusToPiKPi(FemtoFullCollisionMc const& col,
+                         aod::BCsWithTimestamps const&,
+                         FemtoHFMcTracks const& tracks,
+                         aod::McParticles const&,
+                         CandidateDplusMc const& candidates)
+  {
+    // get magnetic field for run
+    getMagneticFieldTesla(col.bc_as<aod::BCsWithTimestamps>());
+
+    fillCharmHadronTable<DecayChannel::DplusToPiKPi, true, false>(col, tracks, candidates);
+  }
+  PROCESS_SWITCH(HfFemtoDreamProducer, processMcDplusToPiKPi, "Provide Mc for DplusToPiKPi", false);
+
+  void processMcDplusToPiKPiWithML(FemtoFullCollisionMc const& col,
+                               aod::BCsWithTimestamps const&,
+                               FemtoHFMcTracks const& tracks,
+                               aod::McParticles const&,
+                               soa::Join<CandidateDplusMc,
+                                         aod::HfMlDplusToPiKPi> const& candidates)
+  {
+    // get magnetic field for run
+    getMagneticFieldTesla(col.bc_as<aod::BCsWithTimestamps>());
+
+    fillCharmHadronTable<DecayChannel::DplusToPiKPi, true, true>(col, tracks, candidates);
+  }
+  PROCESS_SWITCH(HfFemtoDreamProducer, processMcDplusToPiKPiWithML, "Provide Mc for DplusToPiKPi with ml", false);
+
+  void processMcDplusToPiKPiGen(GeneratedMc const& particles)
+  {
+
+    fillCharmHadMcGen(particles);
+  }
+  PROCESS_SWITCH(HfFemtoDreamProducer, processMcDplusToPiKPiGen, "Provide Mc Generated DplusToPiKPi", false);
+
+  /// LcToPKPi
+  void processDataLcToPKPi(FemtoFullCollision const& col,
                            aod::BCsWithTimestamps const&,
                            FemtoHFTracks const& tracks,
                            soa::Filtered<CandidateLc> const& candidates)
   {
     // get magnetic field for run
     getMagneticFieldTesla(col.bc_as<aod::BCsWithTimestamps>());
 
-    fillCharmHadronTable<false, false>(col, tracks, candidates);
+    fillCharmHadronTable<DecayChannel::LcToPKPi, false, false>(col, tracks, candidates);
   }
-  PROCESS_SWITCH(HfFemtoDreamProducer, processDataCharmHad,
-                 "Provide experimental data for charm hadron femto", false);
+  PROCESS_SWITCH(HfFemtoDreamProducer, processDataLcToPKPi,
+                 "Provide experimental data for lctopkpi femto", false);
 
-  void processDataCharmHadWithML(FemtoFullCollision const& col,
+  void processDataLcToPKPiWithML(FemtoFullCollision const& col,
                                  aod::BCsWithTimestamps const&,
                                  FemtoHFTracks const& tracks,
                                  soa::Filtered<soa::Join<CandidateLc,
@@ -609,12 +713,12 @@ struct HfFemtoDreamProducer {
     // get magnetic field for run
     getMagneticFieldTesla(col.bc_as<aod::BCsWithTimestamps>());
 
-    fillCharmHadronTable<false, true>(col, tracks, candidates);
+    fillCharmHadronTable<DecayChannel::LcToPKPi, false, true>(col, tracks, candidates);
   }
-  PROCESS_SWITCH(HfFemtoDreamProducer, processDataCharmHadWithML,
-                 "Provide experimental data for charm hadron femto with ml", false);
+  PROCESS_SWITCH(HfFemtoDreamProducer, processDataLcToPKPiWithML,
+                 "Provide experimental data for lctopkpi femto with ml", false);
 
-  void processMcCharmHad(FemtoFullCollisionMc const& col,
+  void processMcLcToPKPi(FemtoFullCollisionMc const& col,
                          aod::BCsWithTimestamps const&,
                          FemtoHFMcTracks const& tracks,
                          aod::McParticles const&,
@@ -623,11 +727,11 @@ struct HfFemtoDreamProducer {
     // get magnetic field for run
     getMagneticFieldTesla(col.bc_as<aod::BCsWithTimestamps>());
 
-    fillCharmHadronTable<true, false>(col, tracks, candidates);
+    fillCharmHadronTable<DecayChannel::LcToPKPi,true, false>(col, tracks, candidates);
   }
-  PROCESS_SWITCH(HfFemtoDreamProducer, processMcCharmHad, "Provide Mc for charm hadron", false);
+  PROCESS_SWITCH(HfFemtoDreamProducer, processMcLcToPKPi, "Provide Mc for lctopkpi", false);
 
-  void processMcCharmHadWithML(FemtoFullCollisionMc const& col,
+  void processMcLcToPKPiWithML(FemtoFullCollisionMc const& col,
                                aod::BCsWithTimestamps const&,
                                FemtoHFMcTracks const& tracks,
                                aod::McParticles const&,
@@ -637,16 +741,16 @@ struct HfFemtoDreamProducer {
     // get magnetic field for run
     getMagneticFieldTesla(col.bc_as<aod::BCsWithTimestamps>());
 
-    fillCharmHadronTable<true, true>(col, tracks, candidates);
+    fillCharmHadronTable<DecayChannel::LcToPKPi, true, true>(col, tracks, candidates);
   }
-  PROCESS_SWITCH(HfFemtoDreamProducer, processMcCharmHadWithML, "Provide Mc for charm hadron with ml", false);
+  PROCESS_SWITCH(HfFemtoDreamProducer, processMcLcToPKPiWithML, "Provide Mc for lctopkpi with ml", false);
 
-  void processMcCharmHadGen(GeneratedMc const& particles)
+  void processMcLcToPKPiGen(GeneratedMc const& particles)
   {
 
     fillCharmHadMcGen(particles);
   }
-  PROCESS_SWITCH(HfFemtoDreamProducer, processMcCharmHadGen, "Provide Mc Generated charm hadron", false);
+  PROCESS_SWITCH(HfFemtoDreamProducer, processMcLcToPKPiGen, "Provide Mc Generated lctopkpi", false);
 };
 
 WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)
