[PWGJE] GNN b-jet updated (#14334)

Co-authored-by: Changhwan Choi <changhwan.choi@cern.ch>
Co-authored-by: ALICE Action Bot <alibuild@cern.ch>

Author: 3 people

PWGJE/Core/JetTaggingUtilities.h

@@ -1097,11 +1097,11 @@ void analyzeJetTrackInfo4MLnoSV(AnalysisJet const& analysisJet, AnyTracks const&
 
 // Looping over the track info and putting them in the input vector (for GNN b-jet tagging)
 template <typename AnalysisJet, typename AnyTracks, typename AnyOriginalTracks>
-void analyzeJetTrackInfo4GNN(AnalysisJet const& analysisJet, AnyTracks const& /*allTracks*/, AnyOriginalTracks const& /*origTracks*/, std::vector<std::vector<float>>& tracksParams, float trackPtMin = 0.5, int64_t nMaxConstit = 40)
+void analyzeJetTrackInfo4GNN(AnalysisJet const& analysisJet, AnyTracks const& /*allTracks*/, AnyOriginalTracks const& /*origTracks*/, std::vector<std::vector<float>>& tracksParams, float trackPtMin = 0.5, float trackDcaXYMax = 10.0, float trackDcaZMax = 10.0, int64_t nMaxConstit = 40)
 {
   for (const auto& constituent : analysisJet.template tracks_as<AnyTracks>()) {
 
-    if (constituent.pt() < trackPtMin) {
+    if (constituent.pt() < trackPtMin || !trackAcceptanceWithDca(constituent, trackDcaXYMax, trackDcaZMax)) {
       continue;
     }
 

PWGJE/Core/MlResponseHfTagging.h

@@ -23,6 +23,7 @@
 #include <onnxruntime_c_api.h>
 #include <onnxruntime_cxx_api.h>
 
+#include <cmath>
 #include <cstddef>
 #include <cstdint>
 #include <utility>
@@ -208,7 +209,7 @@ class MlResponseHfTagging : public MlResponse<TypeOutputScore>
   static int replaceNaN(std::vector<T>& vec, T value)
   {
     int numNaN = 0;
-    for (auto& el : vec) {
+    for (auto& el : vec) { // o2-linter: disable=const-ref-in-for-loop
       if (std::isnan(el) || std::isinf(el)) {
         el = value;
         ++numNaN;
@@ -364,14 +365,14 @@ class GNNBjetAllocator : public TensorAllocator
   template <typename T>
   T jetFeatureTransform(T feat, int idx) const
   {
-    return (feat - tfJetMean[idx]) / tfJetStdev[idx];
+    return std::tanh((feat - tfJetMean[idx]) / tfJetStdev[idx]);
   }
 
   // Track feature normalization
   template <typename T>
   T trkFeatureTransform(T feat, int idx) const
   {
-    return (feat - tfTrkMean[idx]) / tfTrkStdev[idx];
+    return std::tanh((feat - tfTrkMean[idx]) / tfTrkStdev[idx]);
   }
 
   // Edge input of GNN (fully-connected graph)

PWGJE/TableProducer/jetTaggerHF.cxx

@@ -615,7 +615,7 @@ struct JetTaggerHFTask {
   {
     for (const auto& jet : jets) {
       std::vector<std::vector<float>> trkFeat;
-      jettaggingutilities::analyzeJetTrackInfo4GNN(jet, tracks, origTracks, trkFeat, trackPtMin, nJetConst);
+      jettaggingutilities::analyzeJetTrackInfo4GNN(jet, tracks, origTracks, trkFeat, trackPtMin, trackDcaXYMax, trackDcaZMax, nJetConst);
 
       std::vector<float> jetFeat{jet.pt(), jet.phi(), jet.eta(), jet.mass()};
 
@@ -625,7 +625,13 @@ struct JetTaggerHFTask {
         tensorAlloc.getGNNInput(jetFeat, trkFeat, feat, gnnInput);
 
         auto modelOutput = bMlResponse.getModelOutput(gnnInput, 0);
-        scoreML[jet.globalIndex()] = jettaggingutilities::getDb(modelOutput, fC);
+        float db = jettaggingutilities::getDb(modelOutput, fC);
+        if (!std::isnan(db)) {
+          scoreML[jet.globalIndex()] = db;
+        } else {
+          scoreML[jet.globalIndex()] = 999.;
+          LOGF(debug, "doprocessAlgorithmGNN, Db is NaN (%d)", jet.globalIndex());
+        }
       } else {
         scoreML[jet.globalIndex()] = -999.;
         LOGF(debug, "doprocessAlgorithmGNN, trkFeat.size() <= 0 (%d)", jet.globalIndex());

PWGJE/Tasks/CMakeLists.txt

@@ -336,7 +336,8 @@ if(FastJet_FOUND)
                         COMPONENT_NAME Analysis)
     o2physics_add_dpl_workflow(bjet-tagging-gnn
                         SOURCES bjetTaggingGnn.cxx
-                        PUBLIC_LINK_LIBRARIES O2::Framework O2Physics::PWGJECore O2Physics::AnalysisCore
+                        PUBLIC_LINK_LIBRARIES O2::Framework O2Physics::PWGJECore O2Physics::AnalysisCore O2Physics::EventFilteringUtils
+                        O2Physics::AnalysisCCDB
                         COMPONENT_NAME Analysis)
     o2physics_add_dpl_workflow(jet-shape
                         SOURCES jetShape.cxx