PWGHF: Implement B0 ML and templatize task (#4661)

* Implement B0 ML and templatize task

* Solve MegaLinter issue

* Solve MegaLinter without running clang format

* Add missing lines to fill B0 ML scores histograms in processDataWithB0Ml

* Apply Fabrizio's comments

---------

Co-authored-by: Alexandre Bigot <abigot@sbgat402.in2p3.fr>

Author: AlexBigO

PWGHF/Core/HfMlResponseB0ToDPi.h

@@ -0,0 +1,221 @@
+// 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 HfMlResponsB0ToDPi.h
+/// \brief Class to compute the ML response for B0 → D∓ π± analysis selections
+/// \author Alexandre Bigot <alexandre.bigot@cern.ch>, IPHC Strasbourg
+
+#ifndef PWGHF_CORE_HFMLRESPONSEB0TODPI_H_
+#define PWGHF_CORE_HFMLRESPONSEB0TODPI_H_
+
+#include <map>
+#include <string>
+#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_B0(FEATURE)                                        \
+  {                                                                 \
+#FEATURE, static_cast < uint8_t>(InputFeaturesB0ToDPi::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_B0_FULL(OBJECT, FEATURE, GETTER)   \
+  case static_cast<uint8_t>(InputFeaturesB0ToDPi::FEATURE): { \
+    inputFeatures.emplace_back(OBJECT.GETTER());              \
+    break;                                                    \
+  }
+
+// 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 GETTER function taking OBJECT in argument
+#define CHECK_AND_FILL_VEC_B0_FUNC(OBJECT, FEATURE, GETTER)   \
+  case static_cast<uint8_t>(InputFeaturesB0ToDPi::FEATURE): { \
+    inputFeatures.emplace_back(GETTER(OBJECT));               \
+    break;                                                    \
+  }
+
+// Specific case of CHECK_AND_FILL_VEC_B0_FULL(OBJECT, FEATURE, GETTER)
+// where OBJECT is named candidate and FEATURE = GETTER
+#define CHECK_AND_FILL_VEC_B0(GETTER)                        \
+  case static_cast<uint8_t>(InputFeaturesB0ToDPi::GETTER): { \
+    inputFeatures.emplace_back(candidate.GETTER());          \
+    break;                                                   \
+  }
+
+namespace o2::pid_tpc_tof_utils
+{
+template <typename T1>
+float getTpcTofNSigmaPi1(const T1& prong1)
+{
+  float defaultNSigma = -999.f; // -999.f is the default value set in TPCPIDResponse.h and PIDTOF.h
+
+  bool hasTpc = prong1.hasTPC();
+  bool hasTof = prong1.hasTOF();
+
+  if (hasTpc && hasTof) {
+    float tpcNSigma = prong1.tpcNSigmaPi();
+    float tofNSigma = prong1.tofNSigmaPi();
+    return sqrt(.5f * tpcNSigma * tpcNSigma + .5f * tofNSigma * tofNSigma);
+  }
+  if (hasTpc) {
+    return abs(prong1.tpcNSigmaPi());
+  }
+  if (hasTof) {
+    return abs(prong1.tofNSigmaPi());
+  }
+  return defaultNSigma;
+}
+} // namespace o2::pid_tpc_tof_utils
+
+namespace o2::analysis
+{
+
+enum class InputFeaturesB0ToDPi : uint8_t {
+  ptProng0 = 0,
+  ptProng1,
+  impactParameter0,
+  impactParameter1,
+  impactParameterProduct,
+  chi2PCA,
+  decayLength,
+  decayLengthXY,
+  decayLengthNormalised,
+  decayLengthXYNormalised,
+  cpa,
+  cpaXY,
+  maxNormalisedDeltaIP,
+  prong0MlScoreBkg,
+  prong0MlScorePrompt,
+  prong0MlScoreNonprompt,
+  tpcNSigmaPi1,
+  tofNSigmaPi1,
+  tpcTofNSigmaPi1
+};
+
+template <typename TypeOutputScore = float>
+class HfMlResponseB0ToDPi : public HfMlResponse<TypeOutputScore>
+{
+ public:
+  /// Default constructor
+  HfMlResponseB0ToDPi() = default;
+  /// Default destructor
+  virtual ~HfMlResponseB0ToDPi() = default;
+
+  /// Method to get the input features vector needed for ML inference
+  /// \param candidate is the B0 candidate
+  /// \param prong1 is the candidate's prong1
+  /// \return inputFeatures vector
+  template <bool withDmesMl, typename T1, typename T2>
+  std::vector<float> getInputFeatures(T1 const& candidate,
+                                      T2 const& prong1)
+  {
+    std::vector<float> inputFeatures;
+
+    for (const auto& idx : MlResponse<TypeOutputScore>::mCachedIndices) {
+      if constexpr (withDmesMl) {
+        switch (idx) {
+          CHECK_AND_FILL_VEC_B0(ptProng0);
+          CHECK_AND_FILL_VEC_B0(ptProng1);
+          CHECK_AND_FILL_VEC_B0(impactParameter0);
+          CHECK_AND_FILL_VEC_B0(impactParameter1);
+          CHECK_AND_FILL_VEC_B0(impactParameterProduct);
+          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);
+          CHECK_AND_FILL_VEC_B0(prong0MlScoreBkg);
+          CHECK_AND_FILL_VEC_B0(prong0MlScorePrompt);
+          CHECK_AND_FILL_VEC_B0(prong0MlScoreNonprompt);
+          // TPC PID variable
+          CHECK_AND_FILL_VEC_B0_FULL(prong1, tpcNSigmaPi1, tpcNSigmaPi);
+          // TOF PID variable
+          CHECK_AND_FILL_VEC_B0_FULL(prong1, tofNSigmaPi1, tofNSigmaPi);
+          // Combined PID variables
+          CHECK_AND_FILL_VEC_B0_FUNC(prong1, tpcTofNSigmaPi1, o2::pid_tpc_tof_utils::getTpcTofNSigmaPi1);
+        }
+      } else {
+        switch (idx) {
+          CHECK_AND_FILL_VEC_B0(ptProng0);
+          CHECK_AND_FILL_VEC_B0(ptProng1);
+          CHECK_AND_FILL_VEC_B0(impactParameter0);
+          CHECK_AND_FILL_VEC_B0(impactParameter1);
+          CHECK_AND_FILL_VEC_B0(impactParameterProduct);
+          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(prong1, tpcNSigmaPi1, tpcNSigmaPi);
+          // TOF PID variable
+          CHECK_AND_FILL_VEC_B0_FULL(prong1, tofNSigmaPi1, tofNSigmaPi);
+          // Combined PID variables
+          CHECK_AND_FILL_VEC_B0_FUNC(prong1, tpcTofNSigmaPi1, o2::pid_tpc_tof_utils::getTpcTofNSigmaPi1);
+        }
+      }
+    }
+
+    return inputFeatures;
+  }
+
+ protected:
+  /// Method to fill the map of available input features
+  void setAvailableInputFeatures()
+  {
+    MlResponse<TypeOutputScore>::mAvailableInputFeatures = {
+      FILL_MAP_B0(ptProng0),
+      FILL_MAP_B0(ptProng1),
+      FILL_MAP_B0(impactParameter0),
+      FILL_MAP_B0(impactParameter1),
+      FILL_MAP_B0(impactParameterProduct),
+      FILL_MAP_B0(chi2PCA),
+      FILL_MAP_B0(decayLength),
+      FILL_MAP_B0(decayLengthXY),
+      FILL_MAP_B0(decayLengthNormalised),
+      FILL_MAP_B0(decayLengthXYNormalised),
+      FILL_MAP_B0(cpa),
+      FILL_MAP_B0(cpaXY),
+      FILL_MAP_B0(maxNormalisedDeltaIP),
+      FILL_MAP_B0(prong0MlScoreBkg),
+      FILL_MAP_B0(prong0MlScorePrompt),
+      FILL_MAP_B0(prong0MlScoreNonprompt),
+      // TPC PID variable
+      FILL_MAP_B0(tpcNSigmaPi1),
+      // TOF PID variable
+      FILL_MAP_B0(tofNSigmaPi1),
+      // Combined PID variable
+      FILL_MAP_B0(tpcTofNSigmaPi1)};
+  }
+};
+
+} // namespace o2::analysis
+
+#undef FILL_MAP_B0
+#undef CHECK_AND_FILL_VEC_B0_FULL
+#undef CHECK_AND_FILL_VEC_B0_FUNC
+#undef CHECK_AND_FILL_VEC_B0
+
+#endif // PWGHF_CORE_HFMLRESPONSEB0TODPI_H_

PWGHF/D2H/TableProducer/CMakeLists.txt

@@ -25,7 +25,7 @@ o2physics_add_dpl_workflow(candidate-creator-bplus-reduced
 
 o2physics_add_dpl_workflow(candidate-selector-b0-to-d-pi-reduced
                     SOURCES candidateSelectorB0ToDPiReduced.cxx
-                    PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore
+                    PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2Physics::MLCore
                     COMPONENT_NAME Analysis)
 
 o2physics_add_dpl_workflow(candidate-selector-bplus-to-d0-pi-reduced

PWGHF/D2H/TableProducer/candidateSelectorB0ToDPiReduced.cxx

@@ -21,6 +21,7 @@
 #include "Common/Core/TrackSelectorPID.h"
 
 #include "PWGHF/Core/HfHelper.h"
+#include "PWGHF/Core/HfMlResponseB0ToDPi.h"
 #include "PWGHF/Core/SelectorCuts.h"
 #include "PWGHF/DataModel/CandidateReconstructionTables.h"
 #include "PWGHF/DataModel/CandidateSelectionTables.h"
@@ -33,6 +34,7 @@ using namespace o2::analysis;
 
 struct HfCandidateSelectorB0ToDPiReduced {
   Produces<aod::HfSelB0ToDPi> hfSelB0ToDPiCandidate; // table defined in CandidateSelectionTables.h
+  Produces<aod::HfMlB0ToDPi> hfMlB0ToDPiCandidate;   // table defined in CandidateSelectionTables.h
 
   Configurable<float> ptCandMin{"ptCandMin", 0., "Lower bound of candidate pT"};
   Configurable<float> ptCandMax{"ptCandMax", 50., "Upper bound of candidate pT"};
@@ -57,17 +59,34 @@ struct HfCandidateSelectorB0ToDPiReduced {
   Configurable<LabeledArray<double>> cutsDmesMl{"cutsDmesMl", {hf_cuts_ml::cuts[0], hf_cuts_ml::nBinsPt, hf_cuts_ml::nCutScores, hf_cuts_ml::labelsPt, hf_cuts_ml::labelsDmesCutScore}, "D-meson ML cuts per pT bin"};
   // QA switch
   Configurable<bool> activateQA{"activateQA", false, "Flag to enable QA histogram"};
-
+  // B0 ML inference
+  Configurable<bool> applyB0Ml{"applyB0Ml", false, "Flag to apply ML selections"};
+  Configurable<std::vector<double>> binsPtB0Ml{"binsPtB0Ml", std::vector<double>{hf_cuts_ml::vecBinsPt}, "pT bin limits for ML application"};
+  Configurable<std::vector<int>> cutDirB0Ml{"cutDirB0Ml", std::vector<int>{hf_cuts_ml::vecCutDir}, "Whether to reject score values greater or smaller than the threshold"};
+  Configurable<LabeledArray<double>> cutsB0Ml{"cutsB0Ml", {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<int8_t> nClassesB0Ml{"nClassesB0Ml", (int8_t)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>{"path_ccdb/BDT_B0/"}, "Paths of models on CCDB"};
+  Configurable<std::vector<std::string>> onnxFileNames{"onnxFileNames", std::vector<std::string>{"ModelHandler_onnx_B0ToDPi.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"};
   // variable that will store the value of selectionFlagD (defined in dataCreatorDplusPiReduced.cxx)
   int mySelectionFlagD = -1;
 
-  HfHelper hfHelper;
-  TrackSelectorPi selectorPion;
+  o2::analysis::HfMlResponseB0ToDPi<float> hfMlResponse;
+  float outputMlNotPreselected = -1.;
+  std::vector<float> outputMl = {};
+  o2::ccdb::CcdbApi ccdbApi;
 
-  HistogramRegistry registry{"registry"};
+  TrackSelectorPi selectorPion;
+  HfHelper hfHelper;
 
   using TracksPion = soa::Join<HfRedTracks, HfRedTracksPid>;
 
+  HistogramRegistry registry{"registry"};
+
   void init(InitContext const& initContext)
   {
     std::array<bool, 2> doprocess{doprocessSelection, doprocessSelectionWithDmesMl};
@@ -101,6 +120,18 @@ struct HfCandidateSelectorB0ToDPiReduced {
         registry.get<TH2>(HIST("hSelections"))->GetXaxis()->SetBinLabel(iBin + 1, labels[iBin].data());
       }
     }
+
+    if (applyB0Ml) {
+      hfMlResponse.configure(binsPtB0Ml, cutsB0Ml, cutDirB0Ml, nClassesB0Ml);
+      if (loadModelsFromCCDB) {
+        ccdbApi.init(ccdbUrl);
+        hfMlResponse.setModelPathsCCDB(onnxFileNames, ccdbApi, modelPathsCCDB, timestampCCDB);
+      } else {
+        hfMlResponse.setModelPathsLocal(onnxFileNames);
+      }
+      hfMlResponse.cacheInputFeaturesIndices(namesInputFeatures);
+      hfMlResponse.init();
+    }
   }
 
   /// Main function to perform B0 candidate creation
@@ -125,6 +156,9 @@ struct HfCandidateSelectorB0ToDPiReduced {
       // check if flagged as B0 → D π
       if (!TESTBIT(hfCandB0.hfflag(), hf_cand_b0::DecayType::B0ToDPi)) {
         hfSelB0ToDPiCandidate(statusB0ToDPi);
+        if (applyB0Ml) {
+          hfMlB0ToDPiCandidate(outputMlNotPreselected);
+        }
         if (activateQA) {
           registry.fill(HIST("hSelections"), 1, ptCandB0);
         }
@@ -139,13 +173,19 @@ struct HfCandidateSelectorB0ToDPiReduced {
       // topological cuts
       if (!hfHelper.selectionB0ToDPiTopol(hfCandB0, cuts, binsPt)) {
         hfSelB0ToDPiCandidate(statusB0ToDPi);
+        if (applyB0Ml) {
+          hfMlB0ToDPiCandidate(outputMlNotPreselected);
+        }
         // LOGF(info, "B0 candidate selection failed at topology selection");
         continue;
       }
 
       if constexpr (withDmesMl) { // we include it in the topological selections
         if (!hfHelper.selectionDmesMlScoresForB(hfCandB0, cutsDmesMl, binsPtDmesMl)) {
           hfSelB0ToDPiCandidate(statusB0ToDPi);
+          if (applyB0Ml) {
+            hfMlB0ToDPiCandidate(outputMlNotPreselected);
+          }
           // LOGF(info, "B0 candidate selection failed at D-meson ML selection");
           continue;
         }
@@ -157,8 +197,8 @@ struct HfCandidateSelectorB0ToDPiReduced {
       }
 
       // track-level PID selection
+      auto trackPi = hfCandB0.template prong1_as<TracksPion>();
       if (usePionPid) {
-        auto trackPi = hfCandB0.template prong1_as<TracksPion>();
         int pidTrackPi{TrackSelectorPID::Status::NotApplicable};
         if (usePionPid == 1) {
           pidTrackPi = selectorPion.statusTpcOrTof(trackPi);
@@ -168,13 +208,33 @@ struct HfCandidateSelectorB0ToDPiReduced {
         if (!hfHelper.selectionB0ToDPiPid(pidTrackPi, acceptPIDNotApplicable.value)) {
           // LOGF(info, "B0 candidate selection failed at PID selection");
           hfSelB0ToDPiCandidate(statusB0ToDPi);
+          if (applyB0Ml) {
+            hfMlB0ToDPiCandidate(outputMlNotPreselected);
+          }
           continue;
         }
         SETBIT(statusB0ToDPi, SelectionStep::RecoPID); // RecoPID = 2 --> statusB0ToDPi = 7
         if (activateQA) {
           registry.fill(HIST("hSelections"), 2 + SelectionStep::RecoPID, ptCandB0);
         }
       }
+
+      if (applyB0Ml) {
+        // B0 ML selections
+        std::vector<float> inputFeatures = hfMlResponse.getInputFeatures<withDmesMl>(hfCandB0, trackPi);
+        bool isSelectedMl = hfMlResponse.isSelectedMl(inputFeatures, ptCandB0, outputMl);
+        hfMlB0ToDPiCandidate(outputMl[1]); // storing ML score for signal class
+
+        if (!isSelectedMl) {
+          hfSelB0ToDPiCandidate(statusB0ToDPi);
+          continue;
+        }
+        SETBIT(statusB0ToDPi, SelectionStep::RecoMl); // RecoML = 3 --> statusB0ToDPi = 15 if usePionPid, 11 otherwise
+        if (activateQA) {
+          registry.fill(HIST("hSelections"), 2 + SelectionStep::RecoMl, ptCandB0);
+        }
+      }
+
       hfSelB0ToDPiCandidate(statusB0ToDPi);
       // LOGF(info, "B0 candidate selection passed all selections");
     }