Implement Fabrizio comments

Author: Marcellocosti

PWGHF/Core/HfHelper.h

@@ -33,6 +33,11 @@
 #include <cmath>
 #include <vector>
 
+template <typename T>
+concept IsB0ToDstarPiChannel = requires(T candidate) {
+  candidate.prongD0Id();
+};
+
 class HfHelper
 {
  public:
@@ -634,6 +639,12 @@ class HfHelper
     return candidate.m(std::array{o2::constants::physics::MassDMinus, o2::constants::physics::MassPiPlus});
   }
 
+  template <IsB0ToDstarPiChannel T>
+  auto invMassB0ToDPi(const T& candidate)
+  {
+    return candidate.m(std::array{o2::constants::physics::MassD0, o2::constants::physics::MassPiPlus, o2::constants::physics::MassPiPlus});
+  }
+
   template <typename T>
   auto cosThetaStarB0(const T& candidate)
   {

PWGHF/Core/HfMlResponseB0ToDPi.h

@@ -18,6 +18,7 @@
 #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"
@@ -75,9 +76,12 @@ namespace o2::analysis
 enum class InputFeaturesB0ToDPi : uint8_t {
   ptProng0 = 0,
   ptProng1,
+  ptProng2,
   impactParameter0,
   impactParameter1,
+  impactParameter2,
   impactParameterProduct,
+  impactParameterProngSqSum,
   chi2PCA,
   decayLength,
   decayLengthXY,
@@ -91,7 +95,13 @@ enum class InputFeaturesB0ToDPi : uint8_t {
   prong0MlScoreNonprompt,
   tpcNSigmaPi1,
   tofNSigmaPi1,
-  tpcTofNSigmaPi1
+  tpcTofNSigmaPi1,
+  tpcNSigmaPiBachPi,
+  tofNSigmaPiBachPi,
+  tpcTofNSigmaPiBachPi,
+  tpcNSigmaPiSoftPi,
+  tofNSigmaPiSoftPi,
+  tpcTofNSigmaPiSoftPi
 };
 
 template <typename TypeOutputScore, bool reduced>
@@ -109,7 +119,7 @@ class HfMlResponseB0ToDPi : public HfMlResponse<TypeOutputScore>
   /// \return inputFeatures vector
   template <bool withDmesMl, typename T1, typename T2>
   std::vector<float> getInputFeatures(T1 const& candidate,
-                                      T2 const& prong1,
+                                      T2 const& prongBachPi,
                                       const std::vector<float>* mlScoresD = nullptr)
   {
     std::vector<float> inputFeatures;
@@ -130,11 +140,76 @@ class HfMlResponseB0ToDPi : public HfMlResponse<TypeOutputScore>
         CHECK_AND_FILL_VEC_B0(cpaXY);
         CHECK_AND_FILL_VEC_B0(maxNormalisedDeltaIP);
         // TPC PID variable
-        CHECK_AND_FILL_VEC_B0_FULL(prong1, tpcNSigmaPi1, tpcNSigmaPi);
+        CHECK_AND_FILL_VEC_B0_FULL(prongBachPi, tpcNSigmaPi1, tpcNSigmaPi);
         // TOF PID variable
-        CHECK_AND_FILL_VEC_B0_FULL(prong1, tofNSigmaPi1, tofNSigmaPi);
+        CHECK_AND_FILL_VEC_B0_FULL(prongBachPi, tofNSigmaPi1, tofNSigmaPi);
         // Combined PID variables
-        CHECK_AND_FILL_VEC_B0_FUNC(prong1, tpcTofNSigmaPi1, o2::pid_tpc_tof_utils::getTpcTofNSigmaPi1);
+        CHECK_AND_FILL_VEC_B0_FUNC(prongBachPi, tpcTofNSigmaPi1, o2::pid_tpc_tof_utils::getTpcTofNSigmaPi1);
+      }
+      if constexpr (withDmesMl) {
+        if constexpr (reduced) {
+          switch (idx) {
+            CHECK_AND_FILL_VEC_B0(prong0MlScoreBkg);
+            CHECK_AND_FILL_VEC_B0(prong0MlScorePrompt);
+            CHECK_AND_FILL_VEC_B0(prong0MlScoreNonprompt);
+          }
+        } else {
+          if (mlScoresD) {
+            switch (idx) {
+              CHECK_AND_FILL_VEC_B0_INDEX(prong0MlScoreBkg, *mlScoresD, 0);
+              CHECK_AND_FILL_VEC_B0_INDEX(prong0MlScorePrompt, *mlScoresD, 1);
+              CHECK_AND_FILL_VEC_B0_INDEX(prong0MlScoreNonprompt, *mlScoresD, 2);
+            }
+          } else {
+            LOG(fatal) << "ML scores of D not provided";
+          }
+        }
+      }
+    }
+
+    return inputFeatures;
+  }
+
+  /// 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 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)
+  {
+    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) {
+        CHECK_AND_FILL_VEC_B0(ptProng0);
+        CHECK_AND_FILL_VEC_B0(ptProng1);
+        CHECK_AND_FILL_VEC_B0(ptProng2);
+        CHECK_AND_FILL_VEC_B0(impactParameter0);
+        CHECK_AND_FILL_VEC_B0(impactParameter1);
+        CHECK_AND_FILL_VEC_B0(impactParameter2);
+        CHECK_AND_FILL_VEC_B0(impactParameterProngSqSum);
+        CHECK_AND_FILL_VEC_B0(chi2PCA);
+        CHECK_AND_FILL_VEC_B0(decayLength);
+        CHECK_AND_FILL_VEC_B0(decayLengthXY);
+        CHECK_AND_FILL_VEC_B0(decayLengthNormalised);
+        CHECK_AND_FILL_VEC_B0(decayLengthXYNormalised);
+        CHECK_AND_FILL_VEC_B0(cpa);
+        CHECK_AND_FILL_VEC_B0(cpaXY);
+        CHECK_AND_FILL_VEC_B0(maxNormalisedDeltaIP);
+        // TPC PID variable
+        CHECK_AND_FILL_VEC_B0_FULL(prongBachPi, tpcNSigmaPiBachPi, tpcNSigmaPi);
+        CHECK_AND_FILL_VEC_B0_FULL(prongSoftPi, tpcNSigmaPiSoftPi, tpcNSigmaPiSoftPi);
+        // TOF PID variable
+        CHECK_AND_FILL_VEC_B0_FULL(prongBachPi, tofNSigmaPiBachPi, tofNSigmaPi);
+        CHECK_AND_FILL_VEC_B0_FULL(prongSoftPi, tofNSigmaPiSoftPi, tofNSigmaPiSoftPi);
+        // Combined PID variables
+        CHECK_AND_FILL_VEC_B0_FUNC(prongBachPi, tpcTofNSigmaPiBachPi, o2::pid_tpc_tof_utils::getTpcTofNSigmaPi1);
+        CHECK_AND_FILL_VEC_B0_FUNC(prongSoftPi, tpcTofNSigmaPiSoftPi, o2::pid_tpc_tof_utils::getTpcTofNSigmaSoftPi);
       }
       if constexpr (withDmesMl) {
         if constexpr (reduced) {
@@ -167,9 +242,12 @@ class HfMlResponseB0ToDPi : public HfMlResponse<TypeOutputScore>
     MlResponse<TypeOutputScore>::mAvailableInputFeatures = {
       FILL_MAP_B0(ptProng0),
       FILL_MAP_B0(ptProng1),
+      FILL_MAP_B0(ptProng2),
       FILL_MAP_B0(impactParameter0),
       FILL_MAP_B0(impactParameter1),
+      FILL_MAP_B0(impactParameter2),
       FILL_MAP_B0(impactParameterProduct),
+      FILL_MAP_B0(impactParameterProngSqSum),
       FILL_MAP_B0(chi2PCA),
       FILL_MAP_B0(decayLength),
       FILL_MAP_B0(decayLengthXY),
@@ -183,10 +261,16 @@ class HfMlResponseB0ToDPi : public HfMlResponse<TypeOutputScore>
       FILL_MAP_B0(prong0MlScoreNonprompt),
       // TPC PID variable
       FILL_MAP_B0(tpcNSigmaPi1),
+      FILL_MAP_B0(tpcNSigmaPiSoftPi),
+      FILL_MAP_B0(tpcNSigmaPiBachPi),
       // TOF PID variable
       FILL_MAP_B0(tofNSigmaPi1),
+      FILL_MAP_B0(tofNSigmaPiSoftPi),
+      FILL_MAP_B0(tofNSigmaPiBachPi),
       // Combined PID variable
-      FILL_MAP_B0(tpcTofNSigmaPi1)};
+      FILL_MAP_B0(tpcTofNSigmaPi1),
+      FILL_MAP_B0(tpcTofNSigmaPiSoftPi),
+      FILL_MAP_B0(tpcTofNSigmaPiBachPi)};
   }
 };
 

PWGHF/D2H/DataModel/ReducedDataModel.h

@@ -293,6 +293,13 @@ DECLARE_SOA_TABLE(HfRedTracksCov, "AOD", "HFREDTRACKCOV", //! Table with track c
                   HFTRACKPARCOV_COLUMNS,
                   o2::soa::Marker<1>);
 
+DECLARE_SOA_TABLE(HfRedTracksMom, "AOD", "HFREDTRACKMOM", //! Table with track momentum information for reduced workflow
+                  soa::Index<>,
+                  hf_track_vars_reduced::Px,
+                  hf_track_vars_reduced::Py,
+                  hf_track_vars_reduced::Pz,
+                  hf_track_vars_reduced::Sign);
+
 // CAREFUL: need to follow convention [Name = Description + 's'] in DECLARE_SOA_TABLE(Name, "AOD", Description)
 // to call DECLARE_SOA_INDEX_COLUMN_FULL later on
 DECLARE_SOA_TABLE(HfRedBach0Bases, "AOD", "HFREDBACH0BASE", //! Table with track information for reduced workflow
@@ -551,44 +558,15 @@ DECLARE_SOA_TABLE(HfRedSoftPiPid, "AOD", "HFREDSOFTPIPID",
                   soa::Index<>,
                   hf_cand_dstar::TPCNSigmaPiSoftPi,
                   hf_cand_dstar::TOFNSigmaPiSoftPi,
-                  hf_cand_dstar::TPCNSigmaKaSoftPi,
-                  hf_cand_dstar::TOFNSigmaKaSoftPi,
                   hf_track_vars_reduced::HasTOF,
                   hf_track_vars_reduced::HasTPC,
-                  hf_cand_dstar::TPCTOFNSigmaPiSoftPi<hf_cand_dstar::TPCNSigmaPiSoftPi, hf_cand_dstar::TOFNSigmaPiSoftPi>,
-                  hf_cand_dstar::TPCTOFNSigmaKaSoftPi<hf_cand_dstar::TPCNSigmaKaSoftPi, hf_cand_dstar::TOFNSigmaKaSoftPi>)
+                  hf_cand_dstar::TPCTOFNSigmaPiSoftPi<hf_cand_dstar::TPCNSigmaPiSoftPi, hf_cand_dstar::TOFNSigmaPiSoftPi>)
 
 namespace hf_track_index_reduced
 {
 DECLARE_SOA_INDEX_COLUMN_FULL(SoftPi, softPi, int, HfRedSoftPiBases, ""); //! ReducedCollision index
 }; // namespace hf_track_index_reduced
 
-// CAREFUL: need to follow convention [Name = Description + 's'] in DECLARE_SOA_TABLE(Name, "AOD", Description)
-// to call DECLARE_SOA_INDEX_COLUMN_FULL later on
-DECLARE_SOA_TABLE(HfRedDStars, "AOD", "HFREDDSTAR", //! Table with 2prong candidate information for reduced workflow
-                  o2::soa::Index<>,
-                  hf_track_index_reduced::Prong0Id, hf_track_index_reduced::Prong1Id, hf_track_index_reduced::Prong2Id,
-                  hf_track_index_reduced::SoftPiId, hf_track_index_reduced::HfRedCollisionId,
-                  HFTRACKPAR_COLUMNS,
-                  hf_cand::XSecondaryVertex, hf_cand::YSecondaryVertex, hf_cand::ZSecondaryVertex,
-                  hf_charm_cand_reduced::InvMassHypo0, hf_charm_cand_reduced::InvMassHypo1,
-                  hf_track_vars_reduced::PtProngMin, hf_track_vars_reduced::AbsEtaProngMin,
-                  hf_track_vars_reduced::ItsNClsProngMin, hf_track_vars_reduced::TpcNClsCrossedRowsProngMin, hf_track_vars_reduced::TpcChi2NClProngMax,
-                  aod::track::Px<aod::track::Signed1Pt, aod::track::Snp, aod::track::Alpha>,
-                  aod::track::Py<aod::track::Signed1Pt, aod::track::Snp, aod::track::Alpha>,
-                  aod::track::Pz<aod::track::Signed1Pt, track::Tgl>,
-                  aod::track::PVector<aod::track::Signed1Pt, aod::track::Snp, aod::track::Alpha, aod::track::Tgl>);
-
-DECLARE_SOA_TABLE(HfRedDStarsCov, "AOD", "HFREDDSTARCOV", //! Table with DStar candidate covariance for reduced workflow
-                  o2::soa::Index<>,
-                  HFTRACKPARCOV_COLUMNS,
-                  o2::soa::Marker<1>);
-
-DECLARE_SOA_TABLE(HfRedDStarsMl, "AOD", "HFREDDSTARML", //! Table with DStar candidate ML scores
-                  hf_charm_cand_reduced::MlScoreBkgMassHypo0,
-                  hf_charm_cand_reduced::MlScorePromptMassHypo0,
-                  hf_charm_cand_reduced::MlScoreNonpromptMassHypo0);
-
 // CAREFUL: need to follow convention [Name = Description + 's'] in DECLARE_SOA_TABLE(Name, "AOD", Description)
 // to call DECLARE_SOA_INDEX_COLUMN_FULL later on
 DECLARE_SOA_TABLE(HfRed3Prongs, "AOD", "HFRED3PRONG", //! Table with 3prong candidate information for reduced workflow
@@ -626,6 +604,17 @@ DECLARE_SOA_TABLE_VERSIONED(HfRed3ProngsMl_001, "AOD", "HFRED3PRONGML", 1, //! T
 
 using HfRed3ProngsMl = HfRed3ProngsMl_001;
 
+DECLARE_SOA_TABLE(HfRedMomDDaugs, "AOD", "HFREDMOMDDAUGS", //! Table with 2prong candidate ML scores
+                  hf_cand::PxProng0,
+                  hf_cand::PyProng0,
+                  hf_cand::PzProng0,
+                  hf_cand::PxProng1,
+                  hf_cand::PyProng1,
+                  hf_cand::PzProng1,
+                  hf_cand::PxProng2,
+                  hf_cand::PyProng2,
+                  hf_cand::PzProng2);
+
 // CAREFUL: need to follow convention [Name = Description + 's'] in DECLARE_SOA_TABLE(Name, "AOD", Description)
 // to call DECLARE_SOA_INDEX_COLUMN_FULL later on
 DECLARE_SOA_TABLE(HfRedJpsis, "AOD", "HFREDJPSI", //! Table with J/Psi candidate information for reduced workflow
@@ -755,19 +744,14 @@ using HfRedPidDau0 = HfRedPidDau0s::iterator;
 using HfRedPidDau1 = HfRedPidDau1s::iterator;
 using HfRedPidDau2 = HfRedPidDau2s::iterator;
 
-DECLARE_SOA_TABLE(HfRedB0SoftPi, "AOD", "HFREDB0SOFTPI", //! Table with track information for reduced workflow
-                  o2::soa::Index<>,
-                  hf_track_index::CollisionId,
-                  hf_cand_dstar::PxSoftPi, hf_cand_dstar::PySoftPi, hf_cand_dstar::PzSoftPi,
-                  hf_cand_dstar::DcaYSoftPi, hf_cand_dstar::SigmaYSoftPi);
-
 // Beauty candidates prongs
 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(ProngDStar, prongDStar, int, HfRedDStars, "");     //! ProngDStar index
-DECLARE_SOA_INDEX_COLUMN_FULL(ProngSoftPi, prongSoftPi, int, HfRedB0SoftPi, ""); //! ProngSoftPi 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_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
@@ -777,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::ProngDStarId, hf_cand_b0_reduced::Prong1Id, hf_cand_b0_reduced::ProngSoftPiId);
+                  hf_cand_b0_reduced::ProngD0Id, hf_cand_b0_reduced::ProngSoftPiId, hf_cand_b0_reduced::ProngBachPiId);
 
 DECLARE_SOA_TABLE(HfRedB0DpMls, "AOD", "HFREDB0DPML", //! Table with ML scores for the D+ daughter
                   hf_cand_b0_reduced::Prong0MlScoreBkg,
@@ -786,7 +770,7 @@ DECLARE_SOA_TABLE(HfRedB0DpMls, "AOD", "HFREDB0DPML", //! Table with ML scores f
                   o2::soa::Marker<1>);
 
 using HfRedCandB0 = soa::Join<HfCandB0Ext, HfRedB0Prongs>;
-using HfRedCandB0DStar = soa::Join<HfCandB0Ext, HfRedB0ProngDStars>;
+using HfRedCandB0DStar = soa::Join<HfCandB0DStExt, HfRedB0ProngDStars>;
 
 namespace hf_cand_bplus_reduced
 {
@@ -906,9 +890,8 @@ DECLARE_SOA_TABLE(HfMcCheckDpPis, "AOD", "HFMCCHECKDPPI", //! Table with reconst
 
 // table with results of reconstruction level MC matching
 DECLARE_SOA_TABLE(HfMcRecRedDStarPis, "AOD", "HFMCRECREDDSTPI", //! Table with reconstructed MC information on DStarPi pairs for reduced workflow
-                  hf_cand_b0_reduced::ProngDStarId,
-                  hf_cand_b0_reduced::Prong1Id,
-                  hf_cand_b0_reduced::ProngSoftPiId,
+                  hf_cand_b0_reduced::ProngD0Id,
+                  hf_cand_b0_reduced::ProngBachPiId,
                   hf_cand_b0::FlagMcMatchRec,
                   hf_cand_b0::FlagWrongCollision,
                   hf_cand_b0::DebugMcRec,