[choich08365]

Dear @fcatalan92, @nzardosh, and all concerned,

Could you please take a look at this PR when you have time? I apologise for the large number of changes, and really appreciate your time and effort in reviewing them.
Thank you in advance.

This update extends the GNN b-jet tagger infrastructure to support its track-origin and vertex prediction outputs, enabling future applications such as selecting tracks from b-hadron decays. The current implementation stores jet-tagging and track-origin prediction scores and provides aligned MC truth labels for track origins and vertex indices. Edge-based vertex prediction storage remains deferred.

Tools/ML/model.h, Tools/ML/MlResponse.h

• Added evalModelFull to the model interface and getModelOutputs to MlResponse, returning one vector per output tensor in the ONNX model's output order.
• The existing getModelOutput exposes only the last output tensor. The new interface provides access to all three GNN outputs—edge predictions, track-origin predictions, and jet-flavour predictions—from a single model evaluation.
• Output values are copied into owning vectors so they remain valid after the ONNX output tensors are released.

PWGJE/DataModel/JetTagging.h

• Added tables that associate selected GNN input tracks with their parent jets, MC truth labels, and prediction scores. In the table names below, * denotes Charged or ChargedMCDetectorLevel.

Table	Availability	Contents
*JetGnnTracks	Data and MCD	One row per selected (jet, track) pair, containing the jet index, the JTracks index, and the track's node index in the GNN input.
*JetGnnTrackIds	Data and MCD	One row per jet, containing an array of indices into *JetGnnTracks, accessible through tracksGnn_as<T>().
ChargedMCDetectorLevelJetTrackOriginDef	MCD only	One MC track-origin label per GNN input track.
ChargedMCDetectorLevelJetVertexIndexDef	MCD only	One MC vertex-cluster index per GNN input track.
*JetGnnTrackOriginTags	Data and MCD	One scalar track-origin discriminant, scoreML, per GNN input track.

• Track-origin labels distinguish background, charm, beauty, primary-vertex, and other secondary-vertex tracks. Vertex indices use -1 for unassigned tracks, 0 for the primary vertex, and positive values for secondary vertices.
• Truth and prediction tables share the row count and ordering of *JetGnnTracks, allowing them to be joined and accessed together through tracksGnn_as<soa::Join<...>>().
• *JetGnnTrackIds preserves one row per input jet, including an empty array for jets without selected GNN tracks.
• Edge-related table declarations remain commented out for future implementation.

PWGJE/TableProducer/jetGnnDefinition.cxx

• Added a task that selects constituent tracks for GNN inputs and fills *JetGnnTracks and *JetGnnTrackIds.
• Preserves the association between model-node order and the original constituent tracks, including track replacement when the configured maximum input size is exceeded.
• For MCD jets, also fills ChargedMCDetectorLevelJetTrackOriginDef and ChargedMCDetectorLevelJetVertexIndexDef using vertexClustering, projecting the full constituent labels onto the selected GNN tracks.
• For data jets, fills only the track and index tables because MC truth labels are unavailable.
• Applies collision, jet, and track selection before constructing GNN inputs. Dataset splitting remains the responsibility of downstream tasks.
• Registered jet-gnn-definition in PWGJE/TableProducer/CMakeLists.txt.

PWGJE/TableProducer/jetGnnInference.cxx

• Added a dedicated GNN inference task supporting data and MCD jets:
  • processGnnTag fills the existing jet-tagging tables.
  • processGnnFull additionally fills the track-origin prediction tables using the input-track mapping produced by jet-gnn-definition.
• Includes the jet-tagging functionality provided by processAlgorithmGNN in jetTaggerHF.cxx. A separate task accommodates the additional GNN outputs and their table dependencies.
• Full inference retrieves all output tensors in one model evaluation and computes jet and track discriminants. Edge outputs are currently unused.
• Skips inference for rejected collisions or jets and inputs with fewer than two selected tracks, storing an invalid jet score of -999 while preserving table alignment.
• Added inclusive track-origin score QA and MC truth-separated distributions. MC track-origin labels are used for QA only.
• Retained optional output support for future vertex prediction tables.

PWGJE/Tasks/bjetTreeCreator.cxx

• Updated GNN tree production to read track-origin and vertex labels from the tables produced by jet-gnn-definition, replacing its local vertexClustering call.
• Uses the selected GNN track mapping to associate each stored track with its labels.
• Separated original-track quality variables into processMCJetsForGNNExtrb. The standard processMCJetsForGNN uses derived track tables without subscribing to the original aod::Tracks tables.
• Aligned collision selection with the definition and inference tasks while retaining downstream dataset splitting and sampling.
• Kept these changes confined to GNN tree production.

PWGJE/Tasks/bjetTaggingGnn.cxx

• Introduced task-local collision selection consistent with the GNN definition, inference, and tree-production tasks.
• Applied it to the relevant data and MCD analysis processes, preserving dataset splitting and the dedicated selection-counter processes.

PWGJE/Core/JetTaggingUtilities.h

• Extended analyzeJetTrackInfo4GNN with an optional constituent-index mapping that follows the exact model-node order after track selection and truncation.

• Added getTrackDb as a provisional discriminant for b-decay track selection. Using softmax probabilities for the five track-origin categories, it computes:

  log(Pbeauty /
      [trackFC × Pcharm
       + (1 − trackFC) × (Pbackground + PPV + PotherSV)])
  

[github-actions]

O2 linter results: ❌ 0 errors, ⚠️ 120 warnings, 🔕 5 disabled

[alibuild]

Error while checking build/O2Physics/o2 for ba59bea at 2026-11-10 16:45:

## sw/BUILD/O2Physics-latest/log
/sw/SOURCES/O2Physics/slc9_x86-64-slc9_x86-64/0/PWGJE/DataModel/JetTagging.h:223:1: error: multi-line comment [-Werror=comment]
/sw/SOURCES/O2Physics/slc9_x86-64-slc9_x86-64/0/PWGJE/DataModel/JetTagging.h:223:1: error: multi-line comment [-Werror=comment]
/sw/SOURCES/O2Physics/slc9_x86-64-slc9_x86-64/0/PWGJE/DataModel/JetTagging.h:223:1: error: multi-line comment [-Werror=comment]
/sw/SOURCES/O2Physics/slc9_x86-64-slc9_x86-64/0/PWGJE/DataModel/JetTagging.h:223:1: error: multi-line comment [-Werror=comment]
/sw/SOURCES/O2Physics/slc9_x86-64-slc9_x86-64/0/PWGJE/DataModel/JetTagging.h:223:1: error: multi-line comment [-Werror=comment]
/sw/SOURCES/O2Physics/slc9_x86-64-slc9_x86-64/0/PWGJE/DataModel/JetTagging.h:223:1: error: multi-line comment [-Werror=comment]
/sw/SOURCES/O2Physics/slc9_x86-64-slc9_x86-64/0/PWGJE/DataModel/JetTagging.h:223:1: error: multi-line comment [-Werror=comment]
/sw/SOURCES/O2Physics/slc9_x86-64-slc9_x86-64/0/PWGJE/DataModel/JetTagging.h:223:1: error: multi-line comment [-Werror=comment]
/sw/SOURCES/O2Physics/slc9_x86-64-slc9_x86-64/0/PWGJE/DataModel/JetTagging.h:223:1: error: multi-line comment [-Werror=comment]
/sw/SOURCES/O2Physics/slc9_x86-64-slc9_x86-64/0/PWGJE/DataModel/JetTagging.h:223:1: error: multi-line comment [-Werror=comment]
ninja: build stopped: subcommand failed.

Full log here.

[alibuild]

Error while checking build/O2Physics/code-check for ba59bea at 2026-11-10 17:11:

## sw/BUILD/O2Physics-code-check-latest/log
--
========== List of issues found ==========
++ echo 'To reproduce the results below locally, run: '\''aliBuild build O2Physics-code-check -e ALIBUILD_BASE_HASH=upstream/master'\''.'
To reproduce the results below locally, run: 'aliBuild build O2Physics-code-check -e ALIBUILD_BASE_HASH=upstream/master'.
++ echo 'To apply all automatic fixes, add the '\''-e O2PHYSICS_CHECKER_FIX=1'\'' option.'
To apply all automatic fixes, add the '-e O2PHYSICS_CHECKER_FIX=1' option.
++ echo 'See the documentation for more details at: <https://aliceo2group.github.io/analysis-framework/docs/tools/#running-with-o2physics-code-check>.'
See the documentation for more details at: <https://aliceo2group.github.io/analysis-framework/docs/tools/#running-with-o2physics-code-check>.
++ echo 'Found 20 errors and 70 warnings.'
Found 20 errors and 70 warnings.
++ [[ 20 -gt 0 ]]
++ cat /sw/BUILD/c8c1066e86363a0d7184aefb8fa55fa407ea0050/O2Physics-code-check/errors.txt
PWGJE/DataModel/JetTagging.h:33:1: error: included header utility is not used directly [misc-include-cleaner,-warnings-as-errors]
PWGJE/DataModel/JetTagging.h:34:1: error: included header vector is not used directly [misc-include-cleaner,-warnings-as-errors]
PWGJE/DataModel/JetTagging.h:223:79: error: multi-line // comment [clang-diagnostic-comment,-warnings-as-errors]
PWGJE/DataModel/JetTagging.h:224:75: error: multi-line // comment [clang-diagnostic-comment,-warnings-as-errors]
PWGJE/DataModel/JetTagging.h:225:75: error: multi-line // comment [clang-diagnostic-comment,-warnings-as-errors]
PWGJE/DataModel/JetTagging.h:226:74: error: multi-line // comment [clang-diagnostic-comment,-warnings-as-errors]
PWGJE/DataModel/JetTagging.h:227:74: error: multi-line // comment [clang-diagnostic-comment,-warnings-as-errors]
PWGJE/DataModel/JetTagging.h:228:74: error: multi-line // comment [clang-diagnostic-comment,-warnings-as-errors]
PWGJE/DataModel/JetTagging.h:229:73: error: multi-line // comment [clang-diagnostic-comment,-warnings-as-errors]
PWGJE/DataModel/JetTagging.h:230:73: error: multi-line // comment [clang-diagnostic-comment,-warnings-as-errors]
PWGJE/DataModel/JetTagging.h:231:73: error: multi-line // comment [clang-diagnostic-comment,-warnings-as-errors]
PWGJE/DataModel/JetTagging.h:232:73: error: multi-line // comment [clang-diagnostic-comment,-warnings-as-errors]
PWGJE/DataModel/JetTagging.h:233:73: error: multi-line // comment [clang-diagnostic-comment,-warnings-as-errors]
PWGJE/DataModel/JetTagging.h:234:73: error: multi-line // comment [clang-diagnostic-comment,-warnings-as-errors]
PWGJE/DataModel/JetTagging.h:235:73: error: multi-line // comment [clang-diagnostic-comment,-warnings-as-errors]
PWGJE/DataModel/JetTagging.h:236:73: error: multi-line // comment [clang-diagnostic-comment,-warnings-as-errors]
PWGJE/DataModel/JetTagging.h:237:71: error: multi-line // comment [clang-diagnostic-comment,-warnings-as-errors]
Tools/ML/model.h:187:5: error: no header providing "assert" is directly included [misc-include-cleaner,-warnings-as-errors]
Tools/ML/model.h:190:58: error: no header providing "OrtAllocatorType" is directly included [misc-include-cleaner,-warnings-as-errors]
Tools/ML/model.h:190:95: error: no header providing "OrtMemType" is directly included [misc-include-cleaner,-warnings-as-errors]
++ [[ 70 -gt 0 ]]
++ cat /sw/BUILD/c8c1066e86363a0d7184aefb8fa55fa407ea0050/O2Physics-code-check/warnings.txt
PWGJE/Core/JetTaggingUtilities.h:235:5: warning: do not use 'else' after 'return' [readability-else-after-return]
PWGJE/Core/JetTaggingUtilities.h:294:9: warning: do not use 'else' after 'return' [readability-else-after-return]
PWGJE/Core/JetTaggingUtilities.h:331:9: warning: do not use 'else' after 'return' [readability-else-after-return]
PWGJE/Core/JetTaggingUtilities.h:403:11: warning: do not use 'else' after 'return' [readability-else-after-return]
PWGJE/Core/JetTaggingUtilities.h:438:11: warning: do not use 'else' after 'return' [readability-else-after-return]
PWGJE/Core/JetTaggingUtilities.h:449:5: warning: do not use 'else' after 'return' [readability-else-after-return]
PWGJE/Core/JetTaggingUtilities.h:475:11: warning: do not use 'else' after 'return' [readability-else-after-return]
PWGJE/Core/JetTaggingUtilities.h:495:47: warning: statement should be inside braces [readability-braces-around-statements]
PWGJE/Core/JetTaggingUtilities.h:497:45: warning: statement should be inside braces [readability-braces-around-statements]
PWGJE/Core/JetTaggingUtilities.h:508:41: warning: statement should be inside braces [readability-braces-around-statements]
PWGJE/Core/JetTaggingUtilities.h:510:41: warning: statement should be inside braces [readability-braces-around-statements]
PWGJE/Core/JetTaggingUtilities.h:512:58: warning: statement should be inside braces [readability-braces-around-statements]
PWGJE/Core/JetTaggingUtilities.h:514:58: warning: statement should be inside braces [readability-braces-around-statements]
PWGJE/Core/JetTaggingUtilities.h:518:55: warning: statement should be inside braces [readability-braces-around-statements]
[0 more errors; see full log]

Full log here.