Jet GNN inference, track-origin/vertex definition

Author: Changhwan Choi

PWGJE/Core/JetTaggingUtilities.h

@@ -1130,9 +1130,13 @@ void analyzeJetTrackInfo4GNNwExtra(AnalysisJet const& analysisJet, AnyTracks con
 
 // Looping over the track info and putting them in the input vector (for GNN b-jet tagging)
 template <typename AnalysisJet, typename AnyTracks>
-void analyzeJetTrackInfo4GNN(AnalysisJet const& analysisJet, AnyTracks const& /*allTracks*/, std::vector<std::vector<float>>& tracksParams, float trackPtMin = 0.5, float trackDcaXYMax = 10.0, float trackDcaZMax = 10.0, int64_t nMaxConstit = 40)
+void analyzeJetTrackInfo4GNN(AnalysisJet const& analysisJet, AnyTracks const& /*allTracks*/, std::vector<std::vector<float>>& tracksParams, float trackPtMin = 0.5, float trackDcaXYMax = 10.0, float trackDcaZMax = 10.0, int64_t nMaxConstit = 40, std::vector<std::size_t>* constituentIndices = nullptr)
 {
-  for (const auto& constituent : analysisJet.template tracks_as<AnyTracks>()) {
+  // Optionally retain constituent-array positions in exactly the model-node
+  // order, including replacements when the maximum track count is exceeded.
+  auto const tracks = analysisJet.template tracks_as<AnyTracks>();
+  for (std::size_t constituentIndex = 0; constituentIndex < tracks.size(); ++constituentIndex) {
+    auto const& constituent = tracks[constituentIndex];
 
     if (constituent.pt() < trackPtMin || !trackAcceptanceWithDca(constituent, trackDcaXYMax, trackDcaZMax)) {
       continue;
@@ -1142,15 +1146,22 @@ void analyzeJetTrackInfo4GNN(AnalysisJet const& analysisJet, AnyTracks const& /*
 
     if (static_cast<int64_t>(tracksParams.size()) < nMaxConstit) {
       tracksParams.emplace_back(std::vector<float>{constituent.pt(), constituent.phi(), constituent.eta(), static_cast<float>(constituent.sign()), std::abs(constituent.dcaXY()) * sign, constituent.sigmadcaXY(), std::abs(constituent.dcaZ()) * sign, constituent.sigmadcaZ()});
+      if (constituentIndices) {
+        constituentIndices->push_back(constituentIndex);
+      }
     } else {
       // If there are more than nMaxConstit constituents in the jet, select only nMaxConstit constituents with the highest DCA_XY significance.
       size_t minIdx = 0;
       for (size_t i = 0; i < tracksParams.size(); ++i) {
         if (tracksParams[i][4] / tracksParams[i][5] < tracksParams[minIdx][4] / tracksParams[minIdx][5])
           minIdx = i;
       }
-      if (std::abs(constituent.dcaXY()) * sign / constituent.sigmadcaXY() > tracksParams[minIdx][4] / tracksParams[minIdx][5])
+      if (std::abs(constituent.dcaXY()) * sign / constituent.sigmadcaXY() > tracksParams[minIdx][4] / tracksParams[minIdx][5]) {
         tracksParams[minIdx] = std::vector<float>{constituent.pt(), constituent.phi(), constituent.eta(), static_cast<float>(constituent.sign()), std::abs(constituent.dcaXY()) * sign, constituent.sigmadcaXY(), std::abs(constituent.dcaZ()) * sign, constituent.sigmadcaZ()};
+        if (constituentIndices) {
+          (*constituentIndices)[minIdx] = constituentIndex;
+        }
+      }
     }
   }
 }
@@ -1241,6 +1252,33 @@ T getDb(const std::vector<T>& logits, double fC = 0.018)
   return std::log(softmaxLogits[0] / (fC * softmaxLogits[1] + (1. - fC) * softmaxLogits[2]));
 }
 
+// Track-origin discriminant. Logit order: [background, charm, beauty, PV, other SV].
+template <typename T>
+T getTrackDb(const std::vector<T>& trackLogits, double trackFC)
+{
+  constexpr std::size_t NTrackOrigins = 5;
+  if (trackLogits.size() != NTrackOrigins) {
+    LOG(fatal) << "getTrackDb expects 5 track-origin logits, received " << trackLogits.size();
+  }
+
+  std::array<T, NTrackOrigins> softmaxLogits;
+  T maxLogit = *std::max_element(trackLogits.begin(), trackLogits.end());
+  T sumLogit = 0.;
+  for (std::size_t i = 0; i < NTrackOrigins; ++i) {
+    softmaxLogits[i] = std::exp(trackLogits[i] - maxLogit);
+    sumLogit += softmaxLogits[i];
+  }
+  for (auto& probability : softmaxLogits) {
+    probability /= sumLogit;
+  }
+
+  T denominator = trackFC * softmaxLogits[1] + (1. - trackFC) * (softmaxLogits[0] + softmaxLogits[3] + softmaxLogits[4]);
+  if (denominator == 0.) {
+    LOG(debug) << "jettaggingutilities::getTrackDb, zero denominator";
+  }
+  return std::log(softmaxLogits[2] / denominator);
+}
+
 }; // namespace jettaggingutilities
 
 #endif // PWGJE_CORE_JETTAGGINGUTILITIES_H_

PWGJE/DataModel/JetTagging.h

@@ -170,8 +170,8 @@ JETTAGGING_TABLES_DEF(Full, "F");
 // Tables for full inference of GNN-based b-jet tagger (including track origin, vertex predictions)
 namespace jetgnnindex
 {
-DECLARE_SOA_INDEX_COLUMN(JTrack, track);              //! index in JTracks of this GNN input track (not the parent AOD Tracks index)
-DECLARE_SOA_COLUMN(TrackIndex, trackIndex, uint32_t); //! node index in the GNN input after track selection/truncation; need not match the constituent-array position
+DECLARE_SOA_INDEX_COLUMN(JTrack, track);             //! index in JTracks of this GNN input track (not the parent AOD Tracks index)
+DECLARE_SOA_COLUMN(TrackIndex, trackIndex, uint8_t); //! node index (0..255) in the GNN input after track selection/truncation; need not match the constituent-array position
 // DECLARE_SOA_COLUMN(EdgeIndex, edgeIndex, uint8_t);    //! unique in-jet edge index for the (i,j) constituent-track pair, i<j, triangular order i*(2N-i-1)/2+(j-i-1)
 // DECLARE_SOA_DYNAMIC_COLUMN(TrackIds, trackIds, //! inverse of the triangular edge indexing above: for n constituent tracks, recovers the (i,j) in-jet track index pair (i<j) this edge connects
 //                           [](uint8_t edgeIndex, int n) -> std::pair<uint8_t, uint8_t> {
@@ -188,55 +188,54 @@ DECLARE_SOA_COLUMN(TrackIndex, trackIndex, uint32_t); //! node index in the GNN
 
 namespace jetgnnlabel
 {
-DECLARE_SOA_COLUMN(Origin, origin, uint8_t);           //! trkOrigin from vertexClustering() (0: non-physical primary, 1: charm, 2: beauty, 3: primary vertex, 4: other secondary)
-DECLARE_SOA_COLUMN(VertexIndex, vertexIndex, uint8_t); //! trkVtxIndex from vertexClustering() (0: primary vertex, >0: secondary vertex, ranked by distance to PV)
+DECLARE_SOA_COLUMN(Origin, origin, uint8_t);          //! trkOrigin from vertexClustering() (0: non-physical primary, 1: charm, 2: beauty, 3: primary vertex, 4: other secondary)
+DECLARE_SOA_COLUMN(VertexIndex, vertexIndex, int8_t); //! trkVtxIndex from vertexClustering() (-1: unassigned, 0: primary vertex, >0: secondary vertex, ranked by distance to PV)
 // DECLARE_SOA_COLUMN(SameVertex, sameVertex, bool);      //! true if the (i,j) constituent-track pair of this edge share the same in-jet vertex cluster (same trkVtxIndex from vertexClustering())
 } // namespace jetgnnlabel
 
 namespace jetgnnscore
 {
-DECLARE_SOA_COLUMN(ScoreML, scoreML, float);                        //! GNN prediction score (track/node level, or edge level symmetrised e.g. averaged from the N*N GNN edge matrix)
-DECLARE_SOA_COLUMN(OriginScores, originScores, std::vector<float>); //! one score per origin class, in jetgnnlabel::Origin order (0..4); producer defines logits/probabilities
+DECLARE_SOA_COLUMN(ScoreML, scoreML, float); //! GNN discriminant: one track-origin D_b per GnnTracks row
 } // namespace jetgnnscore
 
-// GNNTrack has one row per evaluated (jet, track) pair. Its jet and track indices
+// GnnTracks has one row per evaluated (jet, track) pair. Its jet and track indices
 // identify the pair independently of the ordering in JetConstituents.
-// GNNTrackIndices is a positional jet extension: write one row for every jet in
+// GnnTrackIds is a positional jet extension: write one row for every jet in
 // Jets order, including an empty array for jets without evaluated tracks. Store
-// GNNTrack row indices (e.g. the backbone producer's lastIndex()), not JTrackIds.
-#define DECLARE_JET_GNNTRACKINDEX_TABLE(_jet_type_, _description_)                                                   \
-  namespace _jet_type_##gnntrackindextable                                                                           \
-  {                                                                                                                  \
-    DECLARE_SOA_INDEX_COLUMN(_jet_type_, jet);                                                                       \
-  }                                                                                                                  \
-  DECLARE_SOA_TABLE(_jet_type_##GNNTrack, "AOD", _description_,                                                      \
-                    o2::soa::Index<>,                                                                                \
-                    _jet_type_##gnntrackindextable::_jet_type_##Id,                                                  \
-                    jetgnnindex::JTrackId,                                                                           \
-                    jetgnnindex::TrackIndex);                                                                        \
-  namespace _jet_type_##gnntrackindices                                                                              \
-  {                                                                                                                  \
-    DECLARE_SOA_ARRAY_INDEX_COLUMN_FULL(_jet_type_##GNNTrack, tracksGnn, int32_t, _jet_type_##GNNTrack, "");         \
-  }                                                                                                                  \
-  DECLARE_SOA_TABLE(_jet_type_##GNNTrackIndices, "AOD", _description_ "I",                                           \
-                    _jet_type_##gnntrackindices::_jet_type_##GNNTrackIds);
+// GnnTracks row indices (e.g. the backbone producer's lastIndex()), not JTrackIds.
+#define DECLARE_JET_GNNTRACKINDEX_TABLE(_jet_type_, _description_)                                            \
+  namespace _jet_type_##gnntrackindextable                                                                    \
+  {                                                                                                           \
+    DECLARE_SOA_INDEX_COLUMN(_jet_type_, jet);                                                                \
+  }                                                                                                           \
+  DECLARE_SOA_TABLE(_jet_type_##GnnTracks, "AOD", _description_,                                              \
+                    o2::soa::Index<>,                                                                         \
+                    _jet_type_##gnntrackindextable::_jet_type_##Id,                                           \
+                    jetgnnindex::JTrackId,                                                                    \
+                    jetgnnindex::TrackIndex);                                                                 \
+  namespace _jet_type_##gnntrackindices                                                                       \
+  {                                                                                                           \
+    DECLARE_SOA_ARRAY_INDEX_COLUMN_FULL(_jet_type_##GnnTrack, tracksGnn, int32_t, _jet_type_##GnnTracks, ""); \
+  }                                                                                                           \
+  DECLARE_SOA_TABLE(_jet_type_##GnnTrackIds, "AOD", _description_ "I",                                        \
+                    _jet_type_##gnntrackindices::_jet_type_##GnnTrackIds);
 
 // #define DECLARE_JET_GNNEDGEINDEX_TABLE(_jet_type_, _description_)          \
-//   namespace _jet_type_##gnnedgeindextable                                 \
-//   {                                                                       \
-//     DECLARE_SOA_INDEX_COLUMN(_jet_type_, jet);                           \
-//   }                                                                      \
-//   DECLARE_SOA_TABLE(_jet_type_##GNNEdge, "AOD", _description_,          \
-//                     o2::soa::Index<>,                                   \
-//                     _jet_type_##gnnedgeindextable::_jet_type_##Id,      \
-//                     jetgnnindex::EdgeIndex,                             \
-//                     jetgnnindex::TrackIds<jetgnnindex::EdgeIndex>);     \
-//   namespace _jet_type_##gnnedgeindices                                  \
-//   {                                                                     \
-//     DECLARE_SOA_ARRAY_INDEX_COLUMN(_jet_type_##GNNEdge, edges);         \
-//   }                                                                     \
-//   DECLARE_SOA_TABLE(_jet_type_##GNNEdgeIndices, "AOD", _description_ "I", \
-//                     _jet_type_##gnnedgeindices::_jet_type_##GNNEdgeIds);
+  namespace _jet_type_##gnnedgeindextable                                 \
+  {                                                                       \
+    DECLARE_SOA_INDEX_COLUMN(_jet_type_, jet);                           \
+  }                                                                      \
+  DECLARE_SOA_TABLE(_jet_type_##GnnEdges, "AOD", _description_,          \
+                    o2::soa::Index<>,                                   \
+                    _jet_type_##gnnedgeindextable::_jet_type_##Id,      \
+                    jetgnnindex::EdgeIndex,                             \
+                    jetgnnindex::TrackIds<jetgnnindex::EdgeIndex>);     \
+  namespace _jet_type_##gnnedgeindices                                  \
+  {                                                                     \
+    DECLARE_SOA_ARRAY_INDEX_COLUMN(_jet_type_##GnnEdge, edges);         \
+  }                                                                     \
+  DECLARE_SOA_TABLE(_jet_type_##GnnEdgeIds, "AOD", _description_ "I", \
+                    _jet_type_##gnnedgeindices::_jet_type_##GnnEdgeIds);
 
 // In-jet track/edge index backbone tables (own the in-jet ordering + the back-reference to the original track)
 DECLARE_JET_GNNTRACKINDEX_TABLE(ChargedJet, "CJGNT");
@@ -248,38 +247,27 @@ DECLARE_JET_GNNTRACKINDEX_TABLE(ChargedMCDetectorLevelJet, "CDJGNT");
 // #undef DECLARE_JET_GNNEDGEINDEX_TABLE
 
 // Truth labels from vertexClustering() (MC detector level only). Thin tables: join positionally with
-// ChargedMCDetectorLevelJetGNNTrack / ChargedMCDetectorLevelJetGNNEdge (same producer loop, same row order).
+// ChargedMCDetectorLevelJetGnnTracks / ChargedMCDetectorLevelJetGnnEdges (same producer loop, same row order).
+// TrackOriginDef must have exactly one row per GnnTracks backbone row, aligned
+// with GnnTrackOriginTags after GNN input selection/truncation. It is not a JTracks
+// extension: a JTrack used by multiple jets has a truth row for each backbone row.
+// tracksGnnIds() addresses this shared row position through a join containing
+// the backbone, the origin definition and the prediction scores.
+// VertexIndexDef uses the same GnnTracks row order and keeps the full-jet vertex
+// labels, without renumbering after GNN track selection/truncation.
 DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetTrackOriginDef, "AOD", "CDJTOD", jetgnnlabel::Origin);
-DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetVertexDef, "AOD", "CDJTVD", jetgnnlabel::VertexIndex);
+DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetVertexIndexDef, "AOD", "CDJTVD", jetgnnlabel::VertexIndex);
 // DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetEdgeDef, "AOD", "CDJEDF", jetgnnlabel::SameVertex);
 
 // GNN track-origin outputs, data and MC detector level. Thin tables: write one
-// score vector per GNNTrack row, in the same row order as the backbone table.
-// The vector contains classes for this track, not scores for multiple tracks.
-DECLARE_SOA_TABLE(ChargedJetTrackTags, "AOD", "CJTRKT", jetgnnscore::OriginScores);
-DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetTrackTags, "AOD", "CDJTRKT", jetgnnscore::OriginScores);
+// scalar score per GnnTracks row, in the same row order as the backbone table.
+DECLARE_SOA_TABLE(ChargedJetGnnTrackOriginTags, "AOD", "CJTRKT", jetgnnscore::ScoreML);
+DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetGnnTrackOriginTags, "AOD", "CDJTRKT", jetgnnscore::ScoreML);
 
 // GNN ScoreML outputs, edge level, data and MC detector level. Thin tables: join positionally with
-// ChargedJetGNNEdge / ChargedMCDetectorLevelJetGNNEdge.
-// DECLARE_SOA_TABLE(ChargedJetEdgeTags, "AOD", "CJEDGT", jetgnnscore::ScoreML);
-// DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetEdgeTags, "AOD", "CDJEDGT", jetgnnscore::ScoreML);
-
-// Prediction joins do not require MC truth tables. Subscribe to the prediction
-// join and JTracks in the consuming task to bind both index levels:
-//   for (auto const& prediction : jet.tracksGnn_as<ChargedJetGNNTrackPredictions>()) {
-//     auto track = prediction.track_as<JTracks>();
-//     auto scores = prediction.originScores();
-//   }
-using ChargedJetGNNTrackPredictions = o2::soa::Join<ChargedJetGNNTrack, ChargedJetTrackTags>;
-using ChargedMCDetectorLevelJetGNNTrackPredictions = o2::soa::Join<ChargedMCDetectorLevelJetGNNTrack, ChargedMCDetectorLevelJetTrackTags>;
-using ChargedJetsWithGNNTrackOrigins = o2::soa::Join<ChargedJets, ChargedJetConstituents, ChargedJetGNNTrackIndices>;
-using ChargedMCDetectorLevelJetsWithGNNTrackOrigins = o2::soa::Join<ChargedMCDetectorLevelJets, ChargedMCDetectorLevelJetConstituents, ChargedMCDetectorLevelJetGNNTrackIndices>;
-
-// Convenience joins including truth labels when available (MC detector level).
-using ChargedJetGNNTrackLabels = ChargedJetGNNTrackPredictions;
-// using ChargedJetGNNEdgeLabels = o2::soa::Join<ChargedJetGNNEdge, ChargedJetEdgeTags>;
-using ChargedMCDetectorLevelJetGNNTrackLabels = o2::soa::Join<ChargedMCDetectorLevelJetGNNTrack, ChargedMCDetectorLevelJetTrackOriginDef, ChargedMCDetectorLevelJetVertexDef, ChargedMCDetectorLevelJetTrackTags>;
-// using ChargedMCDetectorLevelJetGNNEdgeLabels = o2::soa::Join<ChargedMCDetectorLevelJetGNNEdge, ChargedMCDetectorLevelJetEdgeDef, ChargedMCDetectorLevelJetEdgeTags>;
+// ChargedJetGnnEdges / ChargedMCDetectorLevelJetGnnEdges.
+// DECLARE_SOA_TABLE(ChargedJetGnnEdgeTags, "AOD", "CJEDGT", jetgnnscore::ScoreML);
+// DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetGnnEdgeTags, "AOD", "CDJEDGT", jetgnnscore::ScoreML);
 
 } // namespace o2::aod
 

PWGJE/TableProducer/CMakeLists.txt

@@ -86,6 +86,11 @@ o2physics_add_dpl_workflow(jet-taggerhf
                     PUBLIC_LINK_LIBRARIES O2::Framework O2Physics::PWGJECore O2Physics::AnalysisCore O2Physics::MLCore
                     COMPONENT_NAME Analysis)
 
+o2physics_add_dpl_workflow(jet-gnn-definition
+                    SOURCES jetGnnDefinition.cxx
+                    PUBLIC_LINK_LIBRARIES O2::Framework O2Physics::PWGJECore O2Physics::AnalysisCore
+                    COMPONENT_NAME Analysis)
+
 o2physics_add_dpl_workflow(jet-gnn-inference
                     SOURCES jetGnnInference.cxx
                     PUBLIC_LINK_LIBRARIES O2::Framework O2Physics::PWGJECore O2Physics::AnalysisCore O2Physics::MLCore