[PWGHF] Add ML selection for Xic0ToXiPiKf (#11388)

Tao-Fang · alibuild · vkucera · web-flow · commit 754dcbfb8a8d · 2025-06-04T01:17:57.000+02:00
Co-authored-by: ALICE Action Bot &lt;alibuild@cern.ch&gt;
Co-authored-by: Vít Kučera &lt;vit.kucera@cern.ch&gt;
diff --git a/PWGHF/Core/HfMlResponseXic0ToXiPiKf.h b/PWGHF/Core/HfMlResponseXic0ToXiPiKf.h
@@ -0,0 +1,144 @@
+// Copyright 2019-2020 CERN and copyright holders of ALICE O2.
+// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
+// All rights not expressly granted are reserved.
+//
+// This software is distributed under the terms of the GNU General Public
+// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
+//
+// In applying this license CERN does not waive the privileges and immunities
+// granted to it by virtue of its status as an Intergovernmental Organization
+// or submit itself to any jurisdiction.
+
+/// \file HfMlResponseXic0ToXiPiKf.h
+/// \brief Class to compute the ML response for Ξc^0 → Ξ∓ π± kf analysis selections
+/// \author Tao Fang <tao.fang@cern.ch>, Central China Normal University
+
+#ifndef PWGHF_CORE_HFMLRESPONSEXIC0TOXIPIKF_H_
+#define PWGHF_CORE_HFMLRESPONSEXIC0TOXIPIKF_H_
+
+#include <vector>
+
+#include "PWGHF/Core/HfMlResponse.h"
+
+// Fill the map of available input features
+// the key is the feature's name (std::string)
+// the value is the corresponding value in EnumInputFeatures
+#define FILL_MAP_XIC0TOXIPIKF(FEATURE)                                 \
+  {                                                                    \
+    #FEATURE, static_cast<uint8_t>(InputFeaturesXic0ToXiPiKf::FEATURE) \
+  }
+
+// Check if the index of mCachedIndices (index associated to a FEATURE)
+// matches the entry in EnumInputFeatures associated to this FEATURE
+// if so, the inputFeatures vector is filled with the FEATURE's value
+// by calling the corresponding GETTER from OBJECT
+#define CHECK_AND_FILL_VEC_XIC0TOXIPIKF_FULL(OBJECT, FEATURE, GETTER) \
+  case static_cast<uint8_t>(InputFeaturesXic0ToXiPiKf::FEATURE): {    \
+    inputFeatures.emplace_back(OBJECT.GETTER());                      \
+    break;                                                            \
+  }
+
+// where OBJECT is named candidate and FEATURE = GETTER
+#define CHECK_AND_FILL_VEC_XIC0TOXIPIKF(GETTER)                   \
+  case static_cast<uint8_t>(InputFeaturesXic0ToXiPiKf::GETTER): { \
+    inputFeatures.emplace_back(candidate.GETTER());               \
+    break;                                                        \
+  }
+
+namespace o2::analysis
+{
+
+enum class InputFeaturesXic0ToXiPiKf : uint8_t {
+  tpcNSigmaPiFromLambda,
+  tpcNSigmaPiFromCasc,
+  tpcNSigmaPiFromCharmBaryon,
+  dcaCascDau,
+  dcaCharmBaryonDau,
+  kfDcaXYPiFromXic,
+  kfDcaXYCascToPv,
+  cascChi2OverNdf,
+  xicChi2OverNdf,
+  cascldl,
+  chi2TopoCascToPv,
+  chi2TopoCascToXic,
+  cosPaCascToXic,
+  decayLenXYCasc
+};
+
+template <typename TypeOutputScore = float>
+class HfMlResponseXic0ToXiPiKf : public HfMlResponse<TypeOutputScore>
+{
+ public:
+  /// Default constructor
+  HfMlResponseXic0ToXiPiKf() = default;
+  /// Default destructor
+  virtual ~HfMlResponseXic0ToXiPiKf() = default;
+
+  /// Method to get the input features vector needed for ML inference
+  /// \param candidate is the Xic0 candidate
+  /// \return inputFeatures vector
+  template <typename T1, typename T2, typename T3>
+  // std::vector<float> getInputFeatures(T1 const& candidate)
+  std::vector<float> getInputFeatures(T1 const& candidate, T2 const& lamProngPi, T2 const& cascProngPi, T3 const& charmBaryonProngPi)
+  {
+    std::vector<float> inputFeatures;
+
+    for (const auto& idx : MlResponse<TypeOutputScore>::mCachedIndices) {
+      switch (idx) {
+        // PID variables
+        CHECK_AND_FILL_VEC_XIC0TOXIPIKF_FULL(lamProngPi, tpcNSigmaPiFromLambda, tpcNSigmaPi);
+        CHECK_AND_FILL_VEC_XIC0TOXIPIKF_FULL(cascProngPi, tpcNSigmaPiFromCasc, tpcNSigmaPi);
+        CHECK_AND_FILL_VEC_XIC0TOXIPIKF_FULL(charmBaryonProngPi, tpcNSigmaPiFromCharmBaryon, tpcNSigmaPi);
+        // DCA
+        CHECK_AND_FILL_VEC_XIC0TOXIPIKF(dcaCascDau);
+        CHECK_AND_FILL_VEC_XIC0TOXIPIKF(dcaCharmBaryonDau);
+        CHECK_AND_FILL_VEC_XIC0TOXIPIKF(kfDcaXYPiFromXic);
+        CHECK_AND_FILL_VEC_XIC0TOXIPIKF(kfDcaXYCascToPv);
+        // Chi2Geo
+        CHECK_AND_FILL_VEC_XIC0TOXIPIKF(cascChi2OverNdf);
+        CHECK_AND_FILL_VEC_XIC0TOXIPIKF(xicChi2OverNdf);
+        // ldl
+        CHECK_AND_FILL_VEC_XIC0TOXIPIKF(cascldl);
+        // Chi2Topo
+        CHECK_AND_FILL_VEC_XIC0TOXIPIKF(chi2TopoCascToPv);
+        CHECK_AND_FILL_VEC_XIC0TOXIPIKF(chi2TopoCascToXic);
+        // CosPa
+        CHECK_AND_FILL_VEC_XIC0TOXIPIKF(cosPaCascToXic);
+        // Decay length
+        CHECK_AND_FILL_VEC_XIC0TOXIPIKF(decayLenXYCasc);
+      }
+    }
+
+    return inputFeatures;
+  }
+
+ protected:
+  /// Method to fill the map of available input features
+  void setAvailableInputFeatures()
+  {
+    MlResponse<TypeOutputScore>::mAvailableInputFeatures = {
+      FILL_MAP_XIC0TOXIPIKF(tpcNSigmaPiFromLambda),
+      FILL_MAP_XIC0TOXIPIKF(tpcNSigmaPiFromCasc),
+      FILL_MAP_XIC0TOXIPIKF(tpcNSigmaPiFromCharmBaryon),
+      FILL_MAP_XIC0TOXIPIKF(dcaCascDau),
+      FILL_MAP_XIC0TOXIPIKF(dcaCharmBaryonDau),
+      FILL_MAP_XIC0TOXIPIKF(kfDcaXYPiFromXic),
+      FILL_MAP_XIC0TOXIPIKF(kfDcaXYCascToPv),
+      FILL_MAP_XIC0TOXIPIKF(cascChi2OverNdf),
+      FILL_MAP_XIC0TOXIPIKF(xicChi2OverNdf),
+      FILL_MAP_XIC0TOXIPIKF(cascldl),
+      FILL_MAP_XIC0TOXIPIKF(chi2TopoCascToPv),
+      FILL_MAP_XIC0TOXIPIKF(chi2TopoCascToXic),
+      FILL_MAP_XIC0TOXIPIKF(cosPaCascToXic),
+      FILL_MAP_XIC0TOXIPIKF(decayLenXYCasc),
+    };
+  }
+};
+
+} // namespace o2::analysis
+
+#undef FILL_MAP_XIC0TOXIPIKF
+#undef CHECK_AND_FILL_VEC_XIC0TOXIPIKF_FULL
+#undef CHECK_AND_FILL_VEC_XIC0TOXIPIKF
+
+#endif // PWGHF_CORE_HFMLRESPONSEXIC0TOXIPIKF_H_
diff --git a/PWGHF/DataModel/CandidateSelectionTables.h b/PWGHF/DataModel/CandidateSelectionTables.h
@@ -350,6 +350,7 @@ DECLARE_SOA_COLUMN(TofNSigmaPiFromLambda, tofNSigmaPiFromLambda, float);
 DECLARE_SOA_COLUMN(TofNSigmaPrFromLambda, tofNSigmaPrFromLambda, float);
 DECLARE_SOA_COLUMN(PidTpcInfoStored, pidTpcInfoStored, int);
 DECLARE_SOA_COLUMN(PidTofInfoStored, pidTofInfoStored, int);
+DECLARE_SOA_COLUMN(MlProbToXiPi, mlProbToXiPi, std::vector<float>);
 } // namespace hf_sel_toxipi
 
 DECLARE_SOA_TABLE(HfSelToXiPi, "AOD", "HFSELTOXIPI",
@@ -366,6 +367,9 @@ DECLARE_SOA_TABLE(HfSelToXiPiKf, "AOD", "HFSELTOXIPIKF",
                   hf_sel_toxipi::TpcNSigmaPiFromCharmBaryon, hf_sel_toxipi::TpcNSigmaPiFromCasc, hf_sel_toxipi::TpcNSigmaPiFromLambda, hf_sel_toxipi::TpcNSigmaPrFromLambda,
                   hf_sel_toxipi::TofNSigmaPiFromCharmBaryon, hf_sel_toxipi::TofNSigmaPiFromCasc, hf_sel_toxipi::TofNSigmaPiFromLambda, hf_sel_toxipi::TofNSigmaPrFromLambda);
 
+DECLARE_SOA_TABLE(HfMlToXiPiKf, "AOD", "HFMLSELTOXIPIKF",
+                  hf_sel_toxipi::MlProbToXiPi);
+
 namespace hf_sel_toomegapi
 {
 DECLARE_SOA_COLUMN(StatusPidLambda, statusPidLambda, bool);
diff --git a/PWGHF/TableProducer/CMakeLists.txt b/PWGHF/TableProducer/CMakeLists.txt
@@ -174,7 +174,7 @@ o2physics_add_dpl_workflow(candidate-selector-omegac0-to-omega-pi
 
 o2physics_add_dpl_workflow(candidate-selector-xic0-to-xi-pi-kf
                     SOURCES candidateSelectorXic0ToXiPiKf.cxx
-                    PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore
+                    PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2Physics::MLCore
                     COMPONENT_NAME Analysis)
 
 o2physics_add_dpl_workflow(candidate-selector-to-xi-pi
diff --git a/PWGHF/TableProducer/candidateSelectorXic0ToXiPiKf.cxx b/PWGHF/TableProducer/candidateSelectorXic0ToXiPiKf.cxx
@@ -12,6 +12,10 @@
 /// \file candidateSelectorXic0ToXiPiKf.cxx
 /// \brief Xic0 → Xi Pi selection task
 /// \author Ran Tu <ran.tu@cern.ch>, Fudan University
+/// \author Tao Fang <tao.fang@cern.ch>, Central China Normal University
+
+#include <string>
+#include <vector>
 
 #include "CommonConstants/PhysicsConstants.h"
 #include "Framework/AnalysisTask.h"
@@ -20,6 +24,9 @@
 #include "Common/Core/TrackSelection.h"
 #include "Common/Core/TrackSelectorPID.h"
 
+#include "PWGHF/Core/HfMlResponseXic0ToXiPiKf.h"
+#include "PWGHF/Core/SelectorCuts.h"
+
 #include "PWGHF/DataModel/CandidateReconstructionTables.h"
 #include "PWGHF/DataModel/CandidateSelectionTables.h"
 #include "PWGHF/Utils/utilsAnalysis.h"
@@ -39,6 +46,7 @@ enum PidInfoStored {
 /// Struct for applying Xic0 -> Xi pi selection cuts
 struct HfCandidateSelectorXic0ToXiPiKf {
   Produces<aod::HfSelToXiPiKf> hfSelToXiPi;
+  Produces<aod::HfMlToXiPiKf> hfMlToXiPi;
 
   // LF analysis selections
   Configurable<double> radiusCascMin{"radiusCascMin", 0.5, "Min cascade radius"};
@@ -116,6 +124,25 @@ struct HfCandidateSelectorXic0ToXiPiKf {
   Configurable<int> nClustersItsInnBarrMin{"nClustersItsInnBarrMin", 1, "Minimum number of ITS clusters in inner barrel requirement for pi <- charm baryon"};
   Configurable<float> itsChi2PerClusterMax{"itsChi2PerClusterMax", 36, "Maximum value of chi2 fit over ITS clusters for pi <- charm baryon"};
 
+  // ML inference
+  Configurable<bool> applyMl{"applyMl", true, "Flag to apply ML selections"};
+  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"};
+  Configurable<int> nClassesMl{"nClassesMl", static_cast<int>(hf_cuts_ml::NCutScores), "Number of classes in ML model"};
+  Configurable<std::vector<std::string>> namesInputFeatures{"namesInputFeatures", std::vector<std::string>{"feature1", "feature2"}, "Names of ML model input features"};
+
+  // CCDB configuration
+  Configurable<std::string> ccdbUrl{"ccdbUrl", "http://alice-ccdb.cern.ch", "url of the ccdb repository"};
+  Configurable<std::vector<std::string>> modelPathsCCDB{"modelPathsCCDB", std::vector<std::string>{"EventFiltering/PWGHF/BDTXic0ToXipiKf"}, "Paths of models on CCDB"};
+  Configurable<std::vector<std::string>> onnxFileNames{"onnxFileNames", std::vector<std::string>{"ModelHandler_onnx_Xic0ToXipiKf.onnx"}, "ONNX file names for each pT bin (if not from CCDB full path)"};
+  Configurable<int64_t> timestampCCDB{"timestampCCDB", -1, "timestamp of the ONNX file for ML model used to query in CCDB"};
+  Configurable<bool> loadModelsFromCCDB{"loadModelsFromCCDB", false, "Flag to enable or disable the loading of models from CCDB"};
+
+  o2::analysis::HfMlResponseXic0ToXiPiKf<float> hfMlResponse;
+  std::vector<float> outputMlXic0ToXiPiKf = {};
+  o2::ccdb::CcdbApi ccdbApi;
+
   TrackSelectorPr selectorProton;
   TrackSelectorPi selectorPion;
 
@@ -176,6 +203,18 @@ struct HfCandidateSelectorXic0ToXiPiKf {
     registry.add("hSelMassCharmBaryon", "hSelMassCharmBaryon;status;entries", {HistType::kTH1D, {axisSel}});
     registry.add("hSelDcaXYToPvV0Daughters", "hSelDcaXYToPvV0Daughters;status;entries", {HistType::kTH1D, {axisSel}});
     registry.add("hSelDcaXYToPvPiFromCasc", "hSelDcaXYToPvPiFromCasc;status;entries", {HistType::kTH1D, {axisSel}});
+
+    if (applyMl) {
+      hfMlResponse.configure(binsPtMl, cutsMl, cutDirMl, nClassesMl);
+      if (loadModelsFromCCDB) {
+        ccdbApi.init(ccdbUrl);
+        hfMlResponse.setModelPathsCCDB(onnxFileNames, ccdbApi, modelPathsCCDB, timestampCCDB);
+      } else {
+        hfMlResponse.setModelPathsLocal(onnxFileNames);
+      }
+      hfMlResponse.cacheInputFeaturesIndices(namesInputFeatures);
+      hfMlResponse.init();
+    }
   }
 
   void process(aod::HfCandToXiPiKf const& candidates,
@@ -186,6 +225,8 @@ struct HfCandidateSelectorXic0ToXiPiKf {
     // looping over charm baryon candidates
     for (const auto& candidate : candidates) {
 
+      auto ptCand = RecoDecay::pt(candidate.pxCharmBaryon(), candidate.pyCharmBaryon());
+
       bool resultSelections = true; // True if the candidate passes all the selections, False otherwise
 
       auto trackV0PosDauId = candidate.posTrackId();                   // positive V0 daughter
@@ -503,6 +544,17 @@ struct HfCandidateSelectorXic0ToXiPiKf {
         registry.fill(HIST("hSelMassCharmBaryon"), 0);
       }
 
+      // ML selections
+      if (applyMl) {
+        bool isSelectedMlXic0 = false;
+        std::vector<float> inputFeaturesXic0 = hfMlResponse.getInputFeatures(candidate, trackPiFromLam, trackPiFromCasc, trackPiFromCharm);
+        isSelectedMlXic0 = hfMlResponse.isSelectedMl(inputFeaturesXic0, ptCand, outputMlXic0ToXiPiKf);
+        if (!isSelectedMlXic0) {
+          continue;
+        }
+        hfMlToXiPi(outputMlXic0ToXiPiKf);
+      }
+
       hfSelToXiPi(statusPidCharmBaryon, statusPidCascade, statusPidLambda, statusInvMassCharmBaryon, statusInvMassCascade, statusInvMassLambda, resultSelections, infoTpcStored, infoTofStored,
                   trackPiFromCharm.tpcNSigmaPi(), trackPiFromCasc.tpcNSigmaPi(), trackPiFromLam.tpcNSigmaPi(), trackPrFromLam.tpcNSigmaPr(),
                   trackPiFromCharm.tofNSigmaPi(), trackPiFromCasc.tofNSigmaPi(), trackPiFromLam.tofNSigmaPi(), trackPrFromLam.tofNSigmaPr());
