[PWGHF] Adding : ML and skimming in double Dplus correlator + other minor change (#18125)

Co-authored-by: Valerio Di Bella <vdibella@sbgpc51.in2p3.fr>
Co-authored-by: Vít Kučera <26327373+vkucera@users.noreply.github.com>

Author: 3 people

PWGHF/HFC/DataModel/ReducedDMesonPairsTables.h

@@ -38,6 +38,8 @@ DECLARE_SOA_TABLE(HfCandDpMcEvs, "AOD", "HFCANDDPMCEV",
 
 namespace full
 {
+DECLARE_SOA_INDEX_COLUMN(HfCandDpFullEv, hfCandDpFullEv);
+DECLARE_SOA_INDEX_COLUMN(HfCandDpMcEv, hfCandDpMcEv);
 DECLARE_SOA_COLUMN(RSecondaryVertex, rSecondaryVertex, float);                     //! Radius of secondary vertex (cm)
 DECLARE_SOA_COLUMN(PtProng0, ptProng0, float);                                     //! Transverse momentum of prong0 (GeV/c)
 DECLARE_SOA_COLUMN(PProng0, pProng0, float);                                       //! Momentum of prong0 (GeV/c)
@@ -57,8 +59,6 @@ DECLARE_SOA_COLUMN(Eta, eta, float);
 DECLARE_SOA_COLUMN(Phi, phi, float);                                               //! Azimuth angle of candidate
 DECLARE_SOA_COLUMN(E, e, float);                                                   //! Energy of candidate (GeV)
 DECLARE_SOA_COLUMN(Centrality, centrality, float);                                 //! Collision centrality
-DECLARE_SOA_INDEX_COLUMN(HfCandDpMcEv, hfCandDpMcEv);                              //! The Mc collision index this MC particles belongs to
-DECLARE_SOA_INDEX_COLUMN(HfCandDpFullEv, hfCandDpFullEv);                          //! The collision index this candidate belongs to
 DECLARE_SOA_COLUMN(NSigTpcPi0, nSigTpcPi0, float);                                 //! TPC Nsigma separation for prong0 with pion mass hypothesis
 DECLARE_SOA_COLUMN(NSigTpcKa0, nSigTpcKa0, float);                                 //! TPC Nsigma separation for prong0 with kaon mass hypothesis
 DECLARE_SOA_COLUMN(NSigTofPi0, nSigTofPi0, float);                                 //! TOF Nsigma separation for prong0 with pion mass hypothesis
@@ -71,12 +71,12 @@ DECLARE_SOA_COLUMN(NSigTpcPi2, nSigTpcPi2, float);
 DECLARE_SOA_COLUMN(NSigTpcKa2, nSigTpcKa2, float);                                 //! TPC Nsigma separation for prong2 with kaon mass hypothesis
 DECLARE_SOA_COLUMN(NSigTofPi2, nSigTofPi2, float);                                 //! TOF Nsigma separation for prong2 with pion mass hypothesis
 DECLARE_SOA_COLUMN(NSigTofKa2, nSigTofKa2, float);                                 //! TOF Nsigma separation for prong2 with kaon mass hypothesis
-DECLARE_SOA_COLUMN(NSigTpcTofPi0, nSigTpcTofPi0, float);                           //! TPC and TOF combined Nsigma separation for prong0 with pion mass hypothesis
-DECLARE_SOA_COLUMN(NSigTpcTofKa0, nSigTpcTofKa0, float);                           //! TPC and TOF combined Nsigma separation for prong0 with kaon mass hypothesis
-DECLARE_SOA_COLUMN(NSigTpcTofPi1, nSigTpcTofPi1, float);                           //! TPC and TOF combined Nsigma separation for prong1 with pion mass hypothesis
-DECLARE_SOA_COLUMN(NSigTpcTofKa1, nSigTpcTofKa1, float);                           //! TPC and TOF combined Nsigma separation for prong1 with kaon mass hypothesis
-DECLARE_SOA_COLUMN(NSigTpcTofPi2, nSigTpcTofPi2, float);                           //! TPC and TOF combined Nsigma separation for prong2 with pion mass hypothesis
-DECLARE_SOA_COLUMN(NSigTpcTofKa2, nSigTpcTofKa2, float);                           //! TPC and TOF combined Nsigma separation for prong2 with kaon mass hypothesis
+DECLARE_SOA_COLUMN(TpcTofNSigmaPi0, tpcTofNSigmaPi0, float);                       //! TPC and TOF combined Nsigma separation for prong0 with pion mass hypothesis
+DECLARE_SOA_COLUMN(TpcTofNSigmaKa0, tpcTofNSigmaKa0, float);                       //! TPC and TOF combined Nsigma separation for prong0 with kaon mass hypothesis
+DECLARE_SOA_COLUMN(TpcTofNSigmaPi1, tpcTofNSigmaPi1, float);                       //! TPC and TOF combined Nsigma separation for prong1 with pion mass hypothesis
+DECLARE_SOA_COLUMN(TpcTofNSigmaKa1, tpcTofNSigmaKa1, float);                       //! TPC and TOF combined Nsigma separation for prong1 with kaon mass hypothesis
+DECLARE_SOA_COLUMN(TpcTofNSigmaPi2, tpcTofNSigmaPi2, float);                       //! TPC and TOF combined Nsigma separation for prong2 with pion mass hypothesis
+DECLARE_SOA_COLUMN(TpcTofNSigmaKa2, tpcTofNSigmaKa2, float);                       //! TPC and TOF combined Nsigma separation for prong2 with kaon mass hypothesis
 DECLARE_SOA_COLUMN(DecayLength, decayLength, float);                               //! Decay length of candidate (cm)
 DECLARE_SOA_COLUMN(DecayLengthXY, decayLengthXY, float);                           //! Transverse decay length of candidate (cm)
 DECLARE_SOA_COLUMN(DecayLengthNormalised, decayLengthNormalised, float);           //! Normalised decay length of candidate
@@ -96,7 +96,7 @@ DECLARE_SOA_TABLE(HfCandDpMls, "AOD", "HFCANDDPML",
                   full::MlScore0,
                   full::MlScore1)
 
-DECLARE_SOA_TABLE(HfCandDpTinys, "AOD", "HFCANDDPTINY",
+DECLARE_SOA_TABLE(HfCandDpTinys, "AOD", "HFCANDDPTINY", o2::soa::Index<>,
                   full::CandidateSelFlag,
                   full::M,
                   full::Pt,
@@ -107,7 +107,7 @@ DECLARE_SOA_TABLE(HfCandDpTinys, "AOD", "HFCANDDPTINY",
                   hf_cand_mc_flag::OriginMcRec,
                   hf_cand_mc_flag::FlagMcDecayChanRec)
 
-DECLARE_SOA_TABLE(HfCandDpLites, "AOD", "HFCANDDPLITE",
+DECLARE_SOA_TABLE(HfCandDpSimples, "AOD", "HFCANDDPSIMPLE", o2::soa::Index<>,
                   hf_cand::Chi2PCA,
                   full::DecayLength,
                   full::DecayLengthXY,
@@ -126,20 +126,20 @@ DECLARE_SOA_TABLE(HfCandDpLites, "AOD", "HFCANDDPLITE",
                   full::NSigTpcKa0,
                   full::NSigTofPi0,
                   full::NSigTofKa0,
-                  full::NSigTpcTofPi0,
-                  full::NSigTpcTofKa0,
+                  full::TpcTofNSigmaPi0,
+                  full::TpcTofNSigmaKa0,
                   full::NSigTpcPi1,
                   full::NSigTpcKa1,
                   full::NSigTofPi1,
                   full::NSigTofKa1,
-                  full::NSigTpcTofPi1,
-                  full::NSigTpcTofKa1,
+                  full::TpcTofNSigmaPi1,
+                  full::TpcTofNSigmaKa1,
                   full::NSigTpcPi2,
                   full::NSigTpcKa2,
                   full::NSigTofPi2,
                   full::NSigTofKa2,
-                  full::NSigTpcTofPi2,
-                  full::NSigTpcTofKa2,
+                  full::TpcTofNSigmaPi2,
+                  full::TpcTofNSigmaKa2,
                   full::CandidateSelFlag,
                   full::M,
                   full::Pt,
@@ -155,7 +155,7 @@ DECLARE_SOA_TABLE(HfCandDpLites, "AOD", "HFCANDDPLITE",
                   hf_cand_mc_flag::OriginMcRec,
                   hf_cand_mc_flag::FlagMcDecayChanRec)
 
-DECLARE_SOA_TABLE(HfCandDpFulls, "AOD", "HFCANDDPFULL",
+DECLARE_SOA_TABLE(HfCandDpFulls, "AOD", "HFCANDDPFULL", o2::soa::Index<>,
                   hf_cand::XSecondaryVertex,
                   hf_cand::YSecondaryVertex,
                   hf_cand::ZSecondaryVertex,
@@ -201,20 +201,20 @@ DECLARE_SOA_TABLE(HfCandDpFulls, "AOD", "HFCANDDPFULL",
                   full::NSigTpcKa0,
                   full::NSigTofPi0,
                   full::NSigTofKa0,
-                  full::NSigTpcTofPi0,
-                  full::NSigTpcTofKa0,
+                  full::TpcTofNSigmaPi0,
+                  full::TpcTofNSigmaKa0,
                   full::NSigTpcPi1,
                   full::NSigTpcKa1,
                   full::NSigTofPi1,
                   full::NSigTofKa1,
-                  full::NSigTpcTofPi1,
-                  full::NSigTpcTofKa1,
+                  full::TpcTofNSigmaPi1,
+                  full::TpcTofNSigmaKa1,
                   full::NSigTpcPi2,
                   full::NSigTpcKa2,
                   full::NSigTofPi2,
                   full::NSigTofKa2,
-                  full::NSigTpcTofPi2,
-                  full::NSigTpcTofKa2,
+                  full::TpcTofNSigmaPi2,
+                  full::TpcTofNSigmaKa2,
                   full::CandidateSelFlag,
                   full::M,
                   full::Pt,
@@ -233,7 +233,7 @@ DECLARE_SOA_TABLE(HfCandDpFulls, "AOD", "HFCANDDPFULL",
                   hf_cand_mc_flag::OriginMcRec,
                   hf_cand_mc_flag::FlagMcDecayChanRec);
 
-DECLARE_SOA_TABLE(HfCandDpMcPs, "AOD", "HFCANDDPMCP",
+DECLARE_SOA_TABLE(HfCandDpMcPs, "AOD", "HFCANDDPMCP", o2::soa::Index<>,
                   full::Pt,
                   full::Eta,
                   full::Phi,

PWGHF/HFC/TableProducer/CMakeLists.txt

@@ -42,7 +42,7 @@ o2physics_add_dpl_workflow(correlator-dplus-dminus
 o2physics_add_dpl_workflow(correlator-dplus-dplus-reduced
                     SOURCES correlatorDplusDplusReduced.cxx
                     REUSE_FROM HfCandidatesPCH
-                    PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2Physics::EventFilteringUtils
+                    PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2Physics::EventFilteringUtils O2Physics::MLCore
                     COMPONENT_NAME Analysis)
 
 o2physics_add_dpl_workflow(correlator-dplus-hadrons

PWGHF/HFC/TableProducer/correlatorDplusDplusReduced.cxx

@@ -20,6 +20,8 @@
 #include "PWGHF/Core/CentralityEstimation.h"
 #include "PWGHF/Core/DecayChannels.h"
 #include "PWGHF/Core/HfHelper.h"
+#include "PWGHF/Core/HfMlResponseDplusToPiKPi.h"
+#include "PWGHF/Core/SelectorCuts.h"
 #include "PWGHF/DataModel/CandidateReconstructionTables.h"
 #include "PWGHF/DataModel/CandidateSelectionTables.h"
 #include "PWGHF/HFC/DataModel/ReducedDMesonPairsTables.h"
@@ -30,28 +32,33 @@
 #include "Common/DataModel/Centrality.h"
 
 #include <CCDB/BasicCCDBManager.h>
+#include <CCDB/CcdbApi.h>
 #include <Framework/ASoA.h>
 #include <Framework/AnalysisDataModel.h>
 #include <Framework/AnalysisHelpers.h>
 #include <Framework/AnalysisTask.h>
+#include <Framework/Array2D.h>
 #include <Framework/Configurable.h>
 #include <Framework/Expressions.h>
 #include <Framework/HistogramRegistry.h>
 #include <Framework/InitContext.h>
 #include <Framework/runDataProcessing.h>
 
 #include <cstdint>
+#include <cstdlib>
+#include <string>
 #include <vector>
 
 using namespace o2;
+using namespace o2::analysis;
 using namespace o2::framework;
 using namespace o2::framework::expressions;
 using namespace o2::hf_centrality;
 
 /// Writes the full information in an output TTree
 struct HfCorrelatorDplusDplusReduced {
   Produces<o2::aod::HfCandDpFulls> rowCandidateFull;
-  Produces<o2::aod::HfCandDpLites> rowCandidateLite;
+  Produces<o2::aod::HfCandDpSimples> rowCandidateLite;
   Produces<o2::aod::HfCandDpTinys> rowCandidateTiny;
   Produces<o2::aod::HfCandDpFullEvs> rowCandidateFullEvents;
   Produces<o2::aod::HfCandDpMls> rowCandidateMl;
@@ -69,6 +76,34 @@ struct HfCorrelatorDplusDplusReduced {
   Configurable<bool> cfgSkimmedProcessing{"cfgSkimmedProcessing", true, "Enables processing of skimmed datasets"};
   Configurable<bool> skipSingleD{"skipSingleD", true, "Skip collisions with one or less D candidates"};
 
+  Configurable<bool> applyMl{"applyMl", false, "Flag to apply ML selections"};
+  Configurable<bool> applySkimming{"applySkimming", false, "Flag to apply Skimming selections"};
+  Configurable<bool> loadModelsFromCCDB{"loadModelsFromCCDB", false, "Flag to enable or disable the loading of models from CCDB"};
+  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::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/BDTDPlus"}, "Paths of models on CCDB"};
+  Configurable<std::vector<std::string>> onnxFileNames{"onnxFileNames", std::vector<std::string>{"ModelHandler_onnx_DPlusToKPiPi.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<std::vector<std::string>> namesInputFeatures{"namesInputFeatures", std::vector<std::string>{"feature1", "feature2"}, "Names of ML model input features"};
+
+  Configurable<std::vector<double>> cutPtSkimming{"cutPtSkimming", {1, 5, 1000}, "pT bin limits for Skimming application"};
+  Configurable<std::vector<double>> massMin{"massMin", {1.7, 1.7}, "Mass minimal for the cut for each pt bin"};
+  Configurable<std::vector<double>> massMax{"massMax", {2.0, 2.1}, "Mass maximal for the cut for each pt bin"};
+  Configurable<std::vector<double>> cosThetaMin{"cosThetaMin", {0.96, 0.98}, "CosTheta minimal for the cut for each pt bin"};
+  Configurable<std::vector<double>> decayLengthMin{"decayLengthMin", {0.02, 0.03}, "DecayLength minimal for the cut for each pt bin"};
+  Configurable<std::vector<double>> nSigmaTpcMax{"nSigmaTpcMax", {3, 3}, "NsigmaTPC maximal for the cut for each pt bin"};
+  Configurable<std::vector<double>> nSigmaTofMax{"nSigmaTofMax", {3, 3}, "NsigmaTOF maximal for the cut for each pt bin"};
+
+  Configurable<std::vector<double>> binsPtSkimming{"binsPtSkimming", {0}, "pT bin limits for Skimming application"};
+
+  HfMlResponseDplusToPiKPi<float> hfMlResponse;
+
+  std::vector<float> outputMl;
+  o2::ccdb::CcdbApi ccdbApi;
+
   HfHelper hfHelper;
 
   Service<o2::ccdb::BasicCCDBManager> ccdb;
@@ -102,6 +137,39 @@ struct HfCorrelatorDplusDplusReduced {
     if (cfgSkimmedProcessing) {
       zorroSummary.setObject(zorro.getZorroSummary());
     }
+
+    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();
+    }
+  }
+
+  bool skimming(auto const& candidate)
+  {
+    if (candidate.pt() < cutPtSkimming.value.at(0) || candidate.pt() > cutPtSkimming.value.at(cutPtSkimming.value.size() - 1)) {
+      return false;
+    }
+    for (size_t i = 0; i < cutPtSkimming.value.size() - 1; i++) {
+      if (candidate.pt() <= cutPtSkimming.value.at(i + 1)) {
+        if (hfHelper.invMassDplusToPiKPi(candidate) < massMin.value.at(i) ||
+            hfHelper.invMassDplusToPiKPi(candidate) > massMax.value.at(i) ||
+            candidate.cpa() < cosThetaMin.value.at(i) ||
+            candidate.decayLength() < decayLengthMin.value.at(i) ||
+            std::abs(candidate.nSigTofKa1()) > nSigmaTofMax.value.at(i) ||
+            std::abs(candidate.nSigTpcKa1()) > nSigmaTpcMax.value.at(i)) {
+          return false;
+        }
+        return true;
+      }
+    }
+    return false;
   }
 
   template <typename T>
@@ -127,14 +195,14 @@ struct HfCorrelatorDplusDplusReduced {
       channelMc = candidate.flagMcDecayChanRec();
     }
 
-    std::vector<float> outputMl = {-999., -999.};
+    std::vector<float> mlProb = {-999., -999.};
     if constexpr (DoMl) {
       for (unsigned int iclass = 0; iclass < classMlIndexes->size(); iclass++) {
-        outputMl[iclass] = candidate.mlProbDplusToPiKPi()[classMlIndexes->at(iclass)];
+        mlProb[iclass] = candidate.mlProbDplusToPiKPi()[classMlIndexes->at(iclass)];
       }
       rowCandidateMl(
-        outputMl[0],
-        outputMl[1]);
+        mlProb[0],
+        mlProb[1]);
     }
 
     float cent{-1.};
@@ -291,7 +359,9 @@ struct HfCorrelatorDplusDplusReduced {
     static int lastRunNumber = -1;
     // reserve memory
     rowCandidateFullEvents.reserve(collisions.size());
-    if (fillCandidateLiteTable) {
+    if (fillCandidateTinyTable) {
+      rowCandidateTiny.reserve(candidates.size());
+    } else if (fillCandidateLiteTable) {
       rowCandidateLite.reserve(candidates.size());
     } else {
       rowCandidateFull.reserve(candidates.size());
@@ -313,13 +383,31 @@ struct HfCorrelatorDplusDplusReduced {
 
       const auto colId = collision.globalIndex();
       auto candidatesInThisCollision = candidates.sliceBy(tracksPerCollision, colId);
-      if (skipSingleD)
+      if (skipSingleD) {
         if (candidatesInThisCollision.size() < 2) // o2-linter: disable=magic-number (number of candidate must be larger than 1)
+        {
           continue;
+        }
+      }
       fillEvent(collision);
       for (const auto& candidate : candidatesInThisCollision) {
         auto prongCandidate = candidate.prong1_as<aod::Tracks>();
         auto candidateSign = -prongCandidate.sign();
+
+        if (applySkimming &&
+            !skimming(candidate)) {
+          continue;
+        }
+
+        if (applyMl) {
+          std::vector<float> inputFeatures = hfMlResponse.getInputFeatures(candidate);
+          bool const isSelectedMl = hfMlResponse.isSelectedMl(inputFeatures,
+                                                              std::abs(candidate.pt()),
+                                                              outputMl);
+          if (!isSelectedMl) {
+            continue;
+          }
+        }
         fillCandidateTable<aod::Collisions>(candidate, rowCandidateFullEvents.lastIndex(), candidateSign);
       }
     }
@@ -332,7 +420,9 @@ struct HfCorrelatorDplusDplusReduced {
   {
     // reserve memory
     rowCandidateFullEvents.reserve(collisions.size());
-    if (fillCandidateLiteTable) {
+    if (fillCandidateTinyTable) {
+      rowCandidateTiny.reserve(candidates.size());
+    } else if (fillCandidateLiteTable) {
       rowCandidateLite.reserve(candidates.size());
     } else {
       rowCandidateFull.reserve(candidates.size());
@@ -341,13 +431,30 @@ struct HfCorrelatorDplusDplusReduced {
     for (const auto& collision : collisions) { // No skimming for MC data. No Zorro !
       const auto colId = collision.globalIndex();
       auto candidatesInThisCollision = candidates.sliceBy(tracksPerCollision, colId);
-      if (skipSingleD)
+      if (skipSingleD) {
         if (candidatesInThisCollision.size() < 2) // o2-linter: disable=magic-number (number of candidate must be larger than 1)
+        {
           continue;
+        }
+      }
       fillEvent(collision);
       for (const auto& candidate : candidatesInThisCollision) {
         auto prongCandidate = candidate.prong1_as<aod::Tracks>();
         auto candidateSign = -prongCandidate.sign();
+
+        if (applySkimming &&
+            !skimming(candidate)) {
+          continue;
+        }
+        if (applyMl) {
+          std::vector<float> inputFeatures = hfMlResponse.getInputFeatures(candidate);
+          bool const isSelectedMl = hfMlResponse.isSelectedMl(inputFeatures,
+                                                              std::abs(candidate.pt()),
+                                                              outputMl);
+          if (!isSelectedMl) {
+            continue;
+          }
+        }
         fillCandidateTable<aod::Collisions, true>(candidate, rowCandidateFullEvents.lastIndex(), candidateSign);
       }
     }
@@ -364,9 +471,12 @@ struct HfCorrelatorDplusDplusReduced {
     for (const auto& mccollision : mccollisions) { // No skimming for MC data. No Zorro !
       const auto colId = mccollision.globalIndex();
       const auto particlesInThisCollision = mcparticles.sliceBy(mcParticlesPerMcCollision, colId);
-      if (skipSingleD)
+      if (skipSingleD) {
         if (particlesInThisCollision.size() < 2) // o2-linter: disable=magic-number (number of candidate must be larger than 1)
+        {
           continue;
+        }
+      }
       rowCandidateMcCollisions(
         mccollision.posX(),
         mccollision.posY(),