[PWGJE] Adding the IPz as a feature for ML (#11088)

Author: hahassan7

PWGJE/Core/JetTaggingUtilities.h

@@ -79,6 +79,8 @@ struct BJetTrackParams {
   double deltaRJetTrack = 0.0;
   double signedIP2D = 0.0;
   double signedIP2DSign = 0.0;
+  double signedIPz = 0.0;
+  double signedIPzSign = 0.0;
   double signedIP3D = 0.0;
   double signedIP3DSign = 0.0;
   double momFraction = 0.0;
@@ -1011,7 +1013,7 @@ void analyzeJetTrackInfo4ML(AnalysisJet const& analysisJet, AnyTracks const& /*a
       }
     }
 
-    tracksParams.emplace_back(BJetTrackParams{constituent.pt(), constituent.eta(), dotProduct, dotProduct / analysisJet.p(), deltaRJetTrack, std::abs(constituent.dcaXY()) * sign, constituent.sigmadcaXY(), std::abs(constituent.dcaXYZ()) * sign, constituent.sigmadcaXYZ(), constituent.p() / analysisJet.p(), rClosestSV});
+    tracksParams.emplace_back(BJetTrackParams{constituent.pt(), constituent.eta(), dotProduct, dotProduct / analysisJet.p(), deltaRJetTrack, std::abs(constituent.dcaXY()) * sign, constituent.sigmadcaXY(), std::abs(constituent.dcaZ()) * sign, constituent.sigmadcaZ(), std::abs(constituent.dcaXYZ()) * sign, constituent.sigmadcaXYZ(), constituent.p() / analysisJet.p(), rClosestSV});
   }
 
   auto compare = [](BJetTrackParams& tr1, BJetTrackParams& tr2) {
@@ -1036,7 +1038,7 @@ void analyzeJetTrackInfo4MLnoSV(AnalysisJet const& analysisJet, AnyTracks const&
     double dotProduct = RecoDecay::dotProd(std::array<float, 3>{analysisJet.px(), analysisJet.py(), analysisJet.pz()}, std::array<float, 3>{constituent.px(), constituent.py(), constituent.pz()});
     int sign = getGeoSign(analysisJet, constituent);
 
-    tracksParams.emplace_back(BJetTrackParams{constituent.pt(), constituent.eta(), dotProduct, dotProduct / analysisJet.p(), deltaRJetTrack, std::abs(constituent.dcaXY()) * sign, constituent.sigmadcaXY(), std::abs(constituent.dcaXYZ()) * sign, constituent.sigmadcaXYZ(), constituent.p() / analysisJet.p(), 0.0});
+    tracksParams.emplace_back(BJetTrackParams{constituent.pt(), constituent.eta(), dotProduct, dotProduct / analysisJet.p(), deltaRJetTrack, std::abs(constituent.dcaXY()) * sign, constituent.sigmadcaXY(), std::abs(constituent.dcaZ()) * sign, constituent.sigmadcaZ(), std::abs(constituent.dcaXYZ()) * sign, constituent.sigmadcaXYZ(), constituent.p() / analysisJet.p(), 0.0});
   }
 
   auto compare = [](BJetTrackParams& tr1, BJetTrackParams& tr2) {

PWGJE/Core/MlResponseHfTagging.h

@@ -71,6 +71,8 @@ enum class InputFeaturesBTag : uint8_t {
   deltaRJetTrack,
   signedIP2D,
   signedIP2DSign,
+  signedIPz,
+  signedIPzSign,
   signedIP3D,
   signedIP3DSign,
   momFraction,
@@ -148,6 +150,8 @@ class MlResponseHfTagging : public MlResponse<TypeOutputScore>
           CHECK_AND_FILL_VEC_BTAG(trackInput, track, deltaRJetTrack)
           CHECK_AND_FILL_VEC_BTAG(trackInput, track, signedIP2D)
           CHECK_AND_FILL_VEC_BTAG(trackInput, track, signedIP2DSign)
+          CHECK_AND_FILL_VEC_BTAG(trackInput, track, signedIPz)
+          CHECK_AND_FILL_VEC_BTAG(trackInput, track, signedIPzSign)
           CHECK_AND_FILL_VEC_BTAG(trackInput, track, signedIP3D)
           CHECK_AND_FILL_VEC_BTAG(trackInput, track, signedIP3DSign)
           CHECK_AND_FILL_VEC_BTAG(trackInput, track, momFraction)
@@ -192,6 +196,23 @@ class MlResponseHfTagging : public MlResponse<TypeOutputScore>
     }
   }
 
+  /// @brief Method to replace NaN and infinity values in a vector with a specified value
+  /// @param vec is the vector to be processed
+  /// @param value is the value to replace NaN values with
+  /// @return the number of NaN values replaced
+  template <typename T>
+  static int replaceNaN(std::vector<T>& vec, T value)
+  {
+    int numNaN = 0;
+    for (auto& el : vec) {
+      if (std::isnan(el) || std::isinf(el)) {
+        el = value;
+        ++numNaN;
+      }
+    }
+    return numNaN;
+  }
+
   /// Method to get the input features vector needed for ML inference in a 2D vector
   /// \param jet is the b-jet candidate
   /// \param tracks is the vector of tracks associated to the jet
@@ -209,6 +230,10 @@ class MlResponseHfTagging : public MlResponse<TypeOutputScore>
 
     std::vector<std::vector<float>> inputFeatures;
 
+    replaceNaN(jetInput, 0.f);
+    replaceNaN(trackInput, 0.f);
+    replaceNaN(svInput, 0.f);
+
     inputFeatures.push_back(jetInput);
     inputFeatures.push_back(trackInput);
     inputFeatures.push_back(svInput);
@@ -237,6 +262,8 @@ class MlResponseHfTagging : public MlResponse<TypeOutputScore>
     inputFeatures.insert(inputFeatures.end(), trackInput.begin(), trackInput.end());
     inputFeatures.insert(inputFeatures.end(), svInput.begin(), svInput.end());
 
+    replaceNaN(inputFeatures, 0.f);
+
     return inputFeatures;
   }
 
@@ -261,6 +288,8 @@ class MlResponseHfTagging : public MlResponse<TypeOutputScore>
       FILL_MAP_BJET(deltaRJetTrack),
       FILL_MAP_BJET(signedIP2D),
       FILL_MAP_BJET(signedIP2DSign),
+      FILL_MAP_BJET(signedIPz),
+      FILL_MAP_BJET(signedIPzSign),
       FILL_MAP_BJET(signedIP3D),
       FILL_MAP_BJET(signedIP3DSign),
       FILL_MAP_BJET(momFraction),

PWGJE/Tasks/bjetTreeCreator.cxx

@@ -82,7 +82,8 @@ DECLARE_SOA_COLUMN(DotProdTrackJetOverJet, trackdotjetoverjet, float); //! The d
 DECLARE_SOA_COLUMN(DeltaRJetTrack, rjettrack, float);                  //! The DR jet-track
 DECLARE_SOA_COLUMN(SignedIP2D, ip2d, float);                           //! The track signed 2D IP
 DECLARE_SOA_COLUMN(SignedIP2DSign, ip2dsigma, float);                  //! The track signed 2D IP significance
-DECLARE_SOA_COLUMN(SignedIP3D, ip3d, float);                           //! The track signed 3D IP
+DECLARE_SOA_COLUMN(SignedIPz, ipz, float);                             //! The track signed z IP
+DECLARE_SOA_COLUMN(SignedIPzSign, ipzsigma, float);                    //! The track signed z IP significance
 DECLARE_SOA_COLUMN(SignedIP3DSign, ip3dsigma, float);                  //! The track signed 3D IP significance
 DECLARE_SOA_COLUMN(MomFraction, momfraction, float);                   //! The track momentum fraction of the jets
 DECLARE_SOA_COLUMN(DeltaRTrackVertex, rtrackvertex, float);            //! DR between the track and the closest SV, to be decided whether to add to or not
@@ -108,7 +109,8 @@ DECLARE_SOA_TABLE(bjetTracksParams, "AOD", "BJETTRACKSPARAM",
                   trackInfo::DeltaRJetTrack,
                   trackInfo::SignedIP2D,
                   trackInfo::SignedIP2DSign,
-                  trackInfo::SignedIP3D,
+                  trackInfo::SignedIPz,
+                  trackInfo::SignedIPzSign,
                   trackInfo::SignedIP3DSign,
                   trackInfo::MomFraction,
                   trackInfo::DeltaRTrackVertex);
@@ -460,7 +462,7 @@ struct BJetTreeCreator {
       }
 
       if (produceTree) {
-        bjetTracksParamsTable(bjetParamsTable.lastIndex() + 1, constituent.pt(), constituent.eta(), dotProduct, dotProduct / analysisJet.p(), deltaRJetTrack, std::abs(constituent.dcaXY()) * sign, constituent.sigmadcaXY(), std::abs(constituent.dcaXYZ()) * sign, constituent.sigmadcaXYZ(), constituent.p() / analysisJet.p(), RClosestSV);
+        bjetTracksParamsTable(bjetParamsTable.lastIndex() + 1, constituent.pt(), constituent.eta(), dotProduct, dotProduct / analysisJet.p(), deltaRJetTrack, std::abs(constituent.dcaXY()) * sign, constituent.sigmadcaXY(), std::abs(constituent.dcaZ()) * sign, constituent.sigmadcaZ(), constituent.sigmadcaXYZ(), constituent.p() / analysisJet.p(), RClosestSV);
       }
       trackIndices.push_back(bjetTracksParamsTable.lastIndex());
     }
@@ -531,7 +533,7 @@ struct BJetTreeCreator {
       }
 
       if (produceTree) {
-        bjetTracksParamsTable(bjetParamsTable.lastIndex() + 1, constituent.pt(), constituent.eta(), dotProduct, dotProduct / analysisJet.p(), deltaRJetTrack, std::abs(constituent.dcaXY()) * sign, constituent.sigmadcaXY(), std::abs(constituent.dcaXYZ()) * sign, constituent.sigmadcaXYZ(), constituent.p() / analysisJet.p(), 0.);
+        bjetTracksParamsTable(bjetParamsTable.lastIndex() + 1, constituent.pt(), constituent.eta(), dotProduct, dotProduct / analysisJet.p(), deltaRJetTrack, std::abs(constituent.dcaXY()) * sign, constituent.sigmadcaXY(), std::abs(constituent.dcaZ()) * sign, constituent.sigmadcaZ(), constituent.sigmadcaXYZ(), constituent.p() / analysisJet.p(), 0.);
       }
       trackIndices.push_back(bjetTracksParamsTable.lastIndex());
     }