[PWGHF] Add D*pi channel in B0 reduced workflow

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

@@ -75,9 +75,12 @@ namespace o2::analysis
 enum class InputFeaturesB0ToDPi : uint8_t {
   ptProng0 = 0,
   ptProng1,
+  ptProng2,
   impactParameter0,
   impactParameter1,
+  impactParameter2,
   impactParameterProduct,
+  impactParameterProngSqSum,
   chi2PCA,
   decayLength,
   decayLengthXY,
@@ -91,7 +94,13 @@ enum class InputFeaturesB0ToDPi : uint8_t {
   prong0MlScoreNonprompt,
   tpcNSigmaPi1,
   tofNSigmaPi1,
-  tpcTofNSigmaPi1
+  tpcTofNSigmaPi1,
+  tpcNSigmaPiBachPi,
+  tofNSigmaPiBachPi,
+  tpcTofNSigmaPiBachPi,
+  tpcNSigmaPiSoftPi,
+  tofNSigmaPiSoftPi,
+  tpcTofNSigmaPiSoftPi
 };
 
 template <typename TypeOutputScore, bool reduced>
@@ -109,7 +118,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 +139,77 @@ 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 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, 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;
+
+    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,24 +293,12 @@ DECLARE_SOA_TABLE(HfRedTracksCov, "AOD", "HFREDTRACKCOV", //! Table with track c
                   HFTRACKPARCOV_COLUMNS,
                   o2::soa::Marker<1>);
 
-DECLARE_SOA_TABLE(HfRedSoftPiBases, "AOD", "HFREDSOFTPIBASE", //! Table with track information for reduced workflow
-                  soa::Index<>,
-                  hf_track_index_reduced::TrackId,
-                  hf_track_index_reduced::HfRedCollisionId,
-                  HFTRACKPAR_COLUMNS,
-                  hf_track_vars_reduced::ItsNCls,
-                  hf_track_vars_reduced::TpcNClsCrossedRows,
-                  hf_track_vars_reduced::TpcChi2NCl,
-                  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>,
-                  o2::soa::Marker<2>);
-
-DECLARE_SOA_TABLE(HfRedSoftPiCov, "AOD", "HFREDSOFTPICOV", //! Table with track covariance information for reduced workflow
+DECLARE_SOA_TABLE(HfRedTracksMom, "AOD", "HFREDTRACKMOM", //! Table with track momentum information for reduced workflow
                   soa::Index<>,
-                  HFTRACKPARCOV_COLUMNS,
-                  o2::soa::Marker<2>);
+                  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
@@ -546,6 +534,39 @@ DECLARE_SOA_TABLE(HfRed2ProngsMl, "AOD", "HFRED2PRONGML", //! Table with 2prong
                   hf_charm_cand_reduced::MlScorePromptMassHypo1,
                   hf_charm_cand_reduced::MlScoreNonpromptMassHypo1);
 
+// 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(HfRedSoftPiBases, "AOD", "HFREDSOFTPIBASE", //! Table with track information for reduced workflow
+                  soa::Index<>,
+                  hf_track_index_reduced::TrackId,
+                  hf_track_index_reduced::HfRedCollisionId,
+                  HFTRACKPAR_COLUMNS,
+                  hf_track_vars_reduced::ItsNCls,
+                  hf_track_vars_reduced::TpcNClsCrossedRows,
+                  hf_track_vars_reduced::TpcChi2NCl,
+                  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(HfRedSoftPiCov, "AOD", "HFREDSOFTPICOV", //! Table with track covariance information for reduced workflow
+                  soa::Index<>,
+                  HFTRACKPARCOV_COLUMNS,
+                  o2::soa::Marker<2>);
+
+DECLARE_SOA_TABLE(HfRedSoftPiPid, "AOD", "HFREDSOFTPIPID",
+                  soa::Index<>,
+                  hf_cand_dstar::TPCNSigmaPiSoftPi,
+                  hf_cand_dstar::TOFNSigmaPiSoftPi,
+                  hf_track_vars_reduced::HasTOF,
+                  hf_track_vars_reduced::HasTPC,
+                  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(HfRed3Prongs, "AOD", "HFRED3PRONG", //! Table with 3prong candidate information for reduced workflow
@@ -583,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
@@ -711,11 +743,15 @@ using HfRedPidDau2s = HfRedPidDau2s_001;
 using HfRedPidDau0 = HfRedPidDau0s::iterator;
 using HfRedPidDau1 = HfRedPidDau1s::iterator;
 using HfRedPidDau2 = HfRedPidDau2s::iterator;
+
 // 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(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
@@ -724,13 +760,17 @@ DECLARE_SOA_COLUMN(Prong0MlScoreNonprompt, prong0MlScoreNonprompt, float); //! N
 DECLARE_SOA_TABLE(HfRedB0Prongs, "AOD", "HFREDB0PRONG", //! Table with B0 daughter indices
                   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::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,
                   hf_cand_b0_reduced::Prong0MlScorePrompt,
                   hf_cand_b0_reduced::Prong0MlScoreNonprompt,
                   o2::soa::Marker<1>);
 
 using HfRedCandB0 = soa::Join<HfCandB0Ext, HfRedB0Prongs>;
+using HfRedCandB0DStar = soa::Join<HfCandB0DStExt, HfRedB0ProngDStars>;
 
 namespace hf_cand_bplus_reduced
 {
@@ -848,6 +888,15 @@ DECLARE_SOA_TABLE(HfMcCheckDpPis, "AOD", "HFMCCHECKDPPI", //! Table with reconst
                   hf_b0_mc::PdgCodeProng3,
                   o2::soa::Marker<1>);
 
+// 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::ProngD0Id,
+                  hf_cand_b0_reduced::ProngBachPiId,
+                  hf_cand_b0::FlagMcMatchRec,
+                  hf_cand_b0::FlagWrongCollision,
+                  hf_cand_b0::DebugMcRec,
+                  hf_b0_mc::PtMother);
+
 // Table with same size as HFCANDB0
 DECLARE_SOA_TABLE(HfMcRecRedB0s, "AOD", "HFMCRECREDB0", //! Reconstruction-level MC information on B0 candidates for reduced workflow
                   hf_cand_b0::FlagMcMatchRec,