Implement Fabrizio and Vit comments

Author: Marcellocosti

PWGHF/Core/HfMlResponseB0ToDPi.h

@@ -18,7 +18,6 @@
 #define PWGHF_CORE_HFMLRESPONSEB0TODPI_H_
 
 #include "PWGHF/Core/HfMlResponse.h"
-#include "PWGHF/D2H/DataModel/ReducedDataModel.h"
 #include "PWGHF/D2H/Utils/utilsRedDataFormat.h"
 
 #include "Tools/ML/MlResponse.h"
@@ -173,16 +172,17 @@ class HfMlResponseB0ToDPi : public HfMlResponse<TypeOutputScore>
   /// Method to get the input features vector needed for ML inference
   /// \param candidate is the B0 candidate
   /// \param prongBachPi is the candidate's bachelor pion prong
+  /// \param prongSoftPi is the candidate's soft pion prong
   /// \param mlScoresD is the vector of ML scores for the D meson (if available)
   /// \note this method is used for B0 → D*∓ π± candidates with D meson ML scores
   /// \return inputFeatures vector
-  template <bool withDmesMl, IsB0ToDstarPiChannel T1, typename T2>
-  std::vector<float> getInputFeatures(T1 const& candidate,
-                                      T2 const& prongBachPi,
-                                      const std::vector<float>* mlScoresD = nullptr)
+  template <bool withDmesMl, typename T1, typename T2, typename T3>
+  std::vector<float> getInputFeaturesDStarPi(T1 const& candidate,
+                                             T2 const& prongBachPi,
+                                             T3 const& prongSoftPi,
+                                             const std::vector<float>* mlScoresD = nullptr)
   {
     std::vector<float> inputFeatures;
-    auto prongSoftPi = candidate.template prongSoftPi_as<o2::soa::Join<o2::aod::HfRedSoftPiBases, o2::aod::HfRedSoftPiCov, o2::aod::HfRedSoftPiPid>>();
 
     for (const auto& idx : MlResponse<TypeOutputScore>::mCachedIndices) {
       switch (idx) {

PWGHF/D2H/DataModel/ReducedDataModel.h

@@ -749,9 +749,9 @@ namespace hf_cand_b0_reduced
 {
 DECLARE_SOA_INDEX_COLUMN_FULL(Prong0, prong0, int, HfRed3Prongs, "_0");    //! Prong0 index
 DECLARE_SOA_INDEX_COLUMN_FULL(Prong1, prong1, int, HfRedTrackBases, "_1"); //! Prong1 index
-DECLARE_SOA_INDEX_COLUMN_FULL(ProngD0, prongD0, int, HfRed2Prongs, "");    //! ProngD0 index
-DECLARE_SOA_INDEX_COLUMN_FULL(ProngSoftPi, prongSoftPi, int, HfRedTrackBases, ""); //! ProngSoftPi index
-DECLARE_SOA_INDEX_COLUMN_FULL(ProngBachPi, prongBachPi, int, HfRedTrackBases, ""); //! ProngBachPi index
+DECLARE_SOA_INDEX_COLUMN_FULL(ProngD0, prongD0, int, HfRed2Prongs, "_0");    //! ProngD0 index
+DECLARE_SOA_INDEX_COLUMN_FULL(ProngBachPi, prongBachPi, int, HfRedTrackBases, "_1"); //! ProngBachPi index
+DECLARE_SOA_INDEX_COLUMN_FULL(ProngSoftPi, prongSoftPi, int, HfRedSoftPiBases, "_2"); //! ProngSoftPi index
 DECLARE_SOA_COLUMN(Prong0MlScoreBkg, prong0MlScoreBkg, float);             //! Bkg ML score of the D daughter
 DECLARE_SOA_COLUMN(Prong0MlScorePrompt, prong0MlScorePrompt, float);       //! Prompt ML score of the D daughter
 DECLARE_SOA_COLUMN(Prong0MlScoreNonprompt, prong0MlScoreNonprompt, float); //! Nonprompt ML score of the D daughter
@@ -761,7 +761,7 @@ DECLARE_SOA_TABLE(HfRedB0Prongs, "AOD", "HFREDB0PRONG", //! Table with B0 daught
                   hf_cand_b0_reduced::Prong0Id, hf_cand_b0_reduced::Prong1Id);
 
 DECLARE_SOA_TABLE(HfRedB0ProngDStars, "AOD", "HFREDB0PRONGDST", //! Table with B0 daughter indices
-                  hf_cand_b0_reduced::ProngD0Id, hf_cand_b0_reduced::ProngSoftPiId, hf_cand_b0_reduced::ProngBachPiId);
+                  hf_cand_b0_reduced::ProngD0Id, hf_cand_b0_reduced::ProngBachPiId, hf_cand_b0_reduced::ProngSoftPiId);
 
 DECLARE_SOA_TABLE(HfRedB0DpMls, "AOD", "HFREDB0DPML", //! Table with ML scores for the D+ daughter
                   hf_cand_b0_reduced::Prong0MlScoreBkg,

PWGHF/D2H/TableProducer/candidateCreatorB0Reduced.cxx

@@ -228,7 +228,7 @@ struct HfCandidateCreatorB0Reduced {
                             dcaD.getY(), dcaSoftPion.getY(), dcaBachPion.getY(),
                             std::sqrt(dcaD.getSigmaY2()), std::sqrt(dcaSoftPion.getSigmaY2()), std::sqrt(dcaBachPion.getSigmaY2()));
 
-        rowCandidateProngsDStar(candD.globalIndex(), trackSoftPion.globalIndex(), trackBachPion.globalIndex());
+        rowCandidateProngsDStar(candD.globalIndex(), trackBachPion.globalIndex(), trackSoftPion.globalIndex());
 
         if constexpr (withDmesMl) {
           rowCandidateDmesMlScores(candD.mlScoreBkgMassHypo0(), candD.mlScorePromptMassHypo0(), candD.mlScoreNonpromptMassHypo0());

PWGHF/D2H/TableProducer/candidateSelectorB0ToDPiReduced.cxx

@@ -156,20 +156,51 @@ struct HfCandidateSelectorB0ToDPiReduced {
     }
   }
 
-  // For D* channel
-  template <IsB0ToDstarPiChannel T>
-  auto getTrackBachPi(const T& candidate)
-  {
+  /// Utility function to retrieve the bach pion track 
+  ///from the B0 candidate in the D*-pi decay channel
+  /// \param candidate is the B0 candidate
+  /// \return bach pion track
+  template <IsB0ToDstarPiChannel T1>
+  auto getTrackBachPi(const T1& candidate) {
     return candidate.template prongBachPi_as<TracksBachPion>();
   }
 
-  // For non-D* channel
-  template <typename T>
-  auto getTrackBachPi(const T& candidate)
-  {
+  /// Method to get the input features vector needed for ML inference
+  /// \param candidate is the B0 candidate
+  /// \param prongBachPi is the candidate's bachelor pion prong
+  /// \note this method is used for B0 → D*- π+ candidates with D meson ML scores
+  template <bool withDmesMl, IsB0ToDstarPiChannel T1, typename T2>
+  auto getMlInputFeatures(const T1& candB0, const T2& prongBachPi) {
+    auto prongSoftPi = candB0.template prongSoftPi_as<TracksSoftPions>();
+    if constexpr (withDmesMl) {
+        return hfMlResponse.getInputFeaturesDStarPi<true>(candB0, prongBachPi, prongSoftPi);
+    } else {
+        return hfMlResponse.getInputFeaturesDStarPi<false>(candB0, prongBachPi, prongSoftPi);
+    }
+  }
+
+  /// Utility function to retrieve the bach pion track 
+  ///from the B0 candidate in the D-pi decay channel
+  /// \param candidate is the B0 candidate
+  /// \return bach pion track
+  template <typename T1>
+  auto getTrackBachPi(const T1& candidate) {
     return candidate.template prong1_as<TracksBachPion>();
   }
 
+  /// Method to get the input features vector needed for ML inference
+  /// \param candB0 is the B0 candidate
+  /// \param prongBachPi is the candidate's bachelor pion prong
+  /// \note this method is used for B0 → D- π+ candidates with D meson ML scores
+  template <bool withDmesMl, typename T1, typename T2>
+  auto getMlInputFeatures(const T1& candB0, const T2& prongBachPi) {
+    if constexpr (withDmesMl) {
+      return hfMlResponse.getInputFeatures<true>(candB0, prongBachPi);
+    } else {
+      return hfMlResponse.getInputFeatures<false>(candB0, prongBachPi);
+    }
+  }
+
   /// Main function to perform B0 candidate creation
   /// \param withDmesMl is the flag to use the table with ML scores for the D- daughter (only possible if present in the derived data)
   /// \param hfCandsB0 B0 candidates
@@ -220,15 +251,15 @@ struct HfCandidateSelectorB0ToDPiReduced {
         registry.fill(HIST("hSelections"), 2 + SelectionStep::RecoTopol, ptCandB0);
       }
 
-      auto trackPiBachPi = getTrackBachPi(hfCandB0);
+      auto trackBachPi = getTrackBachPi(hfCandB0);
       if (pionPidMethod == PidMethod::TpcOrTof || pionPidMethod == PidMethod::TpcAndTof) {
-        int pidTrackPiBachPi{TrackSelectorPID::Status::NotApplicable};
+        int pidTrackBachPi{TrackSelectorPID::Status::NotApplicable};
         if (pionPidMethod == PidMethod::TpcOrTof) {
-          pidTrackPiBachPi = selectorPion.statusTpcOrTof(trackPiBachPi);
+          pidTrackBachPi = selectorPion.statusTpcOrTof(trackBachPi);
         } else if (pionPidMethod == PidMethod::TpcAndTof) {
-          pidTrackPiBachPi = selectorPion.statusTpcAndTof(trackPiBachPi);
+          pidTrackBachPi = selectorPion.statusTpcAndTof(trackBachPi);
         }
-        if (!hfHelper.selectionB0ToDPiPid(pidTrackPiBachPi, acceptPIDNotApplicable.value)) {
+        if (!hfHelper.selectionB0ToDPiPid(pidTrackBachPi, acceptPIDNotApplicable.value)) {
           // LOGF(info, "B0 candidate selection failed at PID selection");
           hfSelB0ToDPiCandidate(statusB0ToDPi);
           if (applyB0Ml) {
@@ -243,7 +274,7 @@ struct HfCandidateSelectorB0ToDPiReduced {
       }
       if (applyB0Ml) {
         // B0 ML selections
-        std::vector<float> inputFeatures = hfMlResponse.getInputFeatures<withDmesMl>(hfCandB0, trackPiBachPi);
+        std::vector<float> inputFeatures = getMlInputFeatures<withDmesMl>(hfCandB0, trackBachPi);
         bool isSelectedMl = hfMlResponse.isSelectedMl(inputFeatures, ptCandB0, outputMl);
         hfMlB0ToDPiCandidate(outputMl[1]); // storing ML score for signal class