From 561c2b47f26e639c752d5ccf1e00d789a9ea7e6a Mon Sep 17 00:00:00 2001 From: Changhwan Choi Date: Mon, 7 Sep 2026 19:52:03 +0900 Subject: [PATCH 1/7] MlResponse for utilising multiple outputs --- PWGJE/TableProducer/jetGnnInference.cxx | 791 ++++++++++++++++++++++++ Tools/ML/MlResponse.h | 21 + Tools/ML/model.h | 73 +++ 3 files changed, 885 insertions(+) create mode 100644 PWGJE/TableProducer/jetGnnInference.cxx diff --git a/PWGJE/TableProducer/jetGnnInference.cxx b/PWGJE/TableProducer/jetGnnInference.cxx new file mode 100644 index 00000000000..3249c6cb96b --- /dev/null +++ b/PWGJE/TableProducer/jetGnnInference.cxx @@ -0,0 +1,791 @@ +// Copyright 2019-2020 CERN and copyright holders of ALICE O2. +// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders. +// All rights not expressly granted are reserved. +// +// This software is distributed under the terms of the GNU General Public +// License v3 (GPL Version 3), copied verbatim in the file "COPYING". +// +// In applying this license CERN does not waive the privileges and immunities +// granted to it by virtue of its status as an Intergovernmental Organization +// or submit itself to any jurisdiction. + +/// \file jetGnnInference.cxx +/// \brief Task to produce tables for GNN inference +/// +/// \author Changhwan Choi + +#include "PWGJE/Core/JetTaggingUtilities.h" +#include "PWGJE/Core/MlResponseHfTagging.h" +#include "PWGJE/DataModel/Jet.h" +#include "PWGJE/DataModel/JetReducedData.h" +#include "PWGJE/DataModel/JetTagging.h" + +#include "Common/DataModel/TrackSelectionTables.h" +#include "Tools/ML/MlResponse.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace o2; +using namespace o2::framework; +using namespace o2::framework::expressions; + +template +struct JetGnnInferenceTask { + static constexpr double DefaultCutsMl[1][2] = {{0.5, 0.5}}; + + Produces taggingTable; + Produces gnnTrackTable; + Produces gnnVertexTable; + + // configuration topological cut for track and sv + Configurable trackDcaXYMax{"trackDcaXYMax", 1, "minimum DCA xy acceptance for tracks [cm]"}; + Configurable trackDcaZMax{"trackDcaZMax", 2, "minimum DCA z acceptance for tracks [cm]"}; + // Configurable prongsigmaLxyMax{"prongsigmaLxyMax", 100, "maximum sigma of decay length of prongs on xy plane"}; + // Configurable prongsigmaLxyzMax{"prongsigmaLxyzMax", 100, "maximum sigma of decay length of prongs on xyz plane"}; + // Configurable prongIPxyMin{"prongIPxyMin", 0.008, "maximum impact paramter of prongs on xy plane [cm]"}; + // Configurable prongIPxyMax{"prongIPxyMax", 1, "minimum impact parmeter of prongs on xy plane [cm]"}; + // Configurable prongChi2PCAMin{"prongChi2PCAMin", 4, "minimum Chi2 PCA of decay length of prongs"}; + // Configurable prongChi2PCAMax{"prongChi2PCAMax", 100, "maximum Chi2 PCA of decay length of prongs"}; + // Configurable svDispersionMax{"svDispersionMax", 1, "maximum dispersion of sv"}; + + // // configuration about IP method + // Configurable useJetProb{"useJetProb", false, "fill table for track counting algorithm"}; + // Configurable trackProbQA{"trackProbQA", false, "fill track probability histograms separately for geometric positive and negative tracks for QA"}; + // Configurable numCount{"numCount", 3, "number of track counting"}; + // Configurable resoFuncMatching{"resoFuncMatching", 0, "matching parameters of resolution function as MC samble (0: custom, 1: custom & inc, 2: MB, 3: MB & inc, 4: JJ, 5: JJ & inc)"}; + // Configurable> pathsCCDBforIPDataparameter{"pathsCCDBforIPDataparameter", std::vector{"Users/l/leehy/LHC24g4/f_inclusive_0"}, "Paths for fitting parameters of resolution functions of data for IP method on CCDB"}; + // Configurable> pathsCCDBforIPIncparameter{"pathsCCDBforIPIncparameter", std::vector{"Users/l/leehy/LHC24g4/f_inclusive_0"}, "Paths for fitting parameters of resolution functions of inclusive for IP method on CCDB"}; + // Configurable> pathsCCDBforIPBeautyparameter{"pathsCCDBforIPBeautyparameter", std::vector{"Users/l/leehy/LHC24g4/f_inclusive_0"}, "Paths for fitting parameters of resolution functions of beauty for IP method on CCDB"}; + // Configurable> pathsCCDBforIPCharmparameter{"pathsCCDBforIPCharmparameter", std::vector{"Users/l/leehy/LHC24g4/f_inclusive_0"}, "Paths for fitting parameters of resolution functions of charm for IP method on CCDB"}; + // Configurable> pathsCCDBforIPLfparameter{"pathsCCDBforIPLfparameter", std::vector{"Users/l/leehy/LHC24g4/f_inclusive_0"}, "Paths for fitting parameters of resolution functions of light flavour for IP method on CCDB"}; + // Configurable usepTcategorize{"usepTcategorize", false, "p_T categorize TF1 function with Inclusive jet"}; + // Configurable> paramsResoFuncData{"paramsResoFuncData", std::vector{-1.0}, "parameters of gaus(0)+expo(3)+expo(5)+expo(7))"}; + // Configurable> paramsResoFuncIncJetMC{"paramsResoFuncIncJetMC", std::vector{-1.0}, "parameters of gaus(0)+expo(3)+expo(5)+expo(7)))"}; + // Configurable> paramsResoFuncCharmJetMC{"paramsResoFuncCharmJetMC", std::vector{-1.0}, "parameters of gaus(0)+expo(3)+expo(5)+expo(7)))"}; + // Configurable> paramsResoFuncBeautyJetMC{"paramsResoFuncBeautyJetMC", std::vector{-1.0}, "parameters of gaus(0)+expo(3)+expo(5)+expo(7)))"}; + // Configurable> paramsResoFuncLfJetMC{"paramsResoFuncLfJetMC", std::vector{-1.0}, "parameters of gaus(0)+expo(3)+expo(5)+expo(7)))"}; + // Configurable minSignImpXYSig{"minSignImpXYSig", -40.0, "minimum of signed impact parameter significance"}; + // Configurable tagPointForIP{"tagPointForIP", 2.5, "tagging working point for IP"}; + // Configurable tagPointForIPxyz{"tagPointForIPxyz", 2.5, "tagging working point for IP xyz"}; + // Configurable timestampCCDBForIP{"timestampCCDBForIP", -1, "timestamp of the resolution function file for IP method used to query in CCDB"}; + // // configuration about SV method + // Configurable tagPointForSV{"tagPointForSV", 40, "tagging working point for SV"}; + // Configurable tagPointForSVxyz{"tagPointForSVxyz", 40, "tagging working point for SV xyz"}; + + // ML configuration + Configurable nJetConst{"nJetConst", 100, "maximum number of jet consistuents to be used for ML evaluation"}; + Configurable trackPtMin{"trackPtMin", 0.15, "minimum track pT"}; + // Configurable svPtMin{"svPtMin", 0.5, "minimum SV pT"}; + + // Configurable svReductionFactor{"svReductionFactor", 1.0, "factor for how many SVs to keep"}; + + Configurable> binsPtMl{"binsPtMl", std::vector{5., 1000.}, "pT bin limits for ML application"}; + Configurable> cutDirMl{"cutDirMl", std::vector{cuts_ml::CutSmaller, cuts_ml::CutNot}, "Whether to reject score values greater or smaller than the threshold"}; + Configurable> cutsMl{"cutsMl", {DefaultCutsMl[0], 1, 2, {"pT bin 0"}, {"score for default b-jet tagging", "uncer 1"}}, "ML selections per pT bin"}; + Configurable nClassesMl{"nClassesMl", 3, "Number of classes in ML model"}; + Configurable> namesInputFeatures{"namesInputFeatures", std::vector{"feature1", "feature2"}, "Names of ML model input features"}; + // Configurable useDb{"useDb", false, "Flag to use DB for ML model instead of the score"}; + + Configurable ccdbUrl{"ccdbUrl", "http://alice-ccdb.cern.ch", "url of the ccdb repository"}; + Configurable> modelPathsCCDB{"modelPathsCCDB", std::vector{"Users/c/chchoi"}, "Paths of models on CCDB"}; + Configurable> onnxFileNames{"onnxFileNames", std::vector{"gnnbjettagger_LHC25a2b_v241.onnx"}, "ONNX file names for each pT bin (if not from CCDB full path)"}; + Configurable timestampCCDB{"timestampCCDB", -1, "timestamp of the ONNX file for ML model used to query in CCDB"}; + Configurable loadModelsFromCCDB{"loadModelsFromCCDB", false, "Flag to enable or disable the loading of models from CCDB"}; + + // GNN configuration + Configurable jetpTMin{"jetpTMin", 5., "minimum jet pT"}; + Configurable dbMin{"dbMin", -10., "minimum GNN Db"}; + Configurable dbMax{"dbMax", 20., "maximum GNN Db"}; + Configurable fC{"fC", 0.018, "Parameter f_c for D_b calculation"}; + Configurable nJetFeat{"nJetFeat", 4, "Number of jet GNN input features"}; + Configurable nTrkFeat{"nTrkFeat", 13, "Number of track GNN input features"}; + Configurable nTrkOrigin{"nTrkOrigin", 5, "Number of track origin categories"}; + Configurable> transformFeatureJetMean{"transformFeatureJetMean", + std::vector{-999}, + "Mean values for each GNN input feature (jet)"}; + Configurable> transformFeatureJetStdev{"transformFeatureJetStdev", + std::vector{-999}, + "Stdev values for each GNN input feature (jet)"}; + Configurable> transformFeatureTrkMean{"transformFeatureTrkMean", + std::vector{-999}, + "Mean values for each GNN input feature (track)"}; + Configurable> transformFeatureTrkStdev{"transformFeatureTrkStdev", + std::vector{-999}, + "Stdev values for each GNN input feature (track)"}; + Configurable tfFuncTypeGNN{"tfFuncTypeGNN", "linear", "Transformation function type for GNN"}; + + // axis spec + ConfigurableAxis binTrackProbability{"binTrackProbability", {100, 0.f, 1.f}, ""}; + ConfigurableAxis binJetFlavour{"binJetFlavour", {6, -0.5, 5.5}, ""}; + + o2::analysis::MlResponseHfTagging bMlResponse; + o2::ccdb::CcdbApi ccdbApi; + + using JetTracksExt = soa::Join; + using OriginalTracks = soa::Join; + + bool useResoFuncFromIncJet = false; + int maxOrder = -1; + int resoFuncMatch = 0; + + std::unique_ptr fSignImpXYSigData = nullptr; + std::unique_ptr fSignImpXYSigIncJetMC = nullptr; + std::unique_ptr fSignImpXYSigCharmJetMC = nullptr; + std::unique_ptr fSignImpXYSigBeautyJetMC = nullptr; + std::unique_ptr fSignImpXYSigLfJetMC = nullptr; + + std::vector> vecParamsDataJetCCDB; + std::vector> vecParamsIncJetMcCCDB; + std::vector> vecParamsBeautyJetMcCCDB; + std::vector> vecParamsCharmJetMcCCDB; + std::vector> vecParamsLfJetMcCCDB; + + std::vector> vecfSignImpXYSigDataJetCCDB; + std::vector> vecfSignImpXYSigIncJetMcCCDB; + std::vector> vecfSignImpXYSigCharmJetMcCCDB; + std::vector> vecfSignImpXYSigBeautyJetMcCCDB; + std::vector> vecfSignImpXYSigLfJetMcCCDB; + + std::vector decisionNonML; + std::vector scoreML; + + o2::analysis::GNNBjetAllocator tensorAlloc; + + template + float calculateJetProbability(int origin, T const& jet, U const& tracks, bool const& isMC = false) + { + float jetProb = -1.0; + if (!isMC) { + if (usepTcategorize) { + jetProb = jettaggingutilities::getJetProbability(vecfSignImpXYSigDataJetCCDB, jet, tracks, trackDcaXYMax, trackDcaZMax, minSignImpXYSig); + } else { + jetProb = jettaggingutilities::getJetProbability(fSignImpXYSigData, jet, tracks, trackDcaXYMax, trackDcaZMax, minSignImpXYSig); + } + } else { + if (useResoFuncFromIncJet) { + if (usepTcategorize) { + jetProb = jettaggingutilities::getJetProbability(vecfSignImpXYSigIncJetMcCCDB, jet, tracks, trackDcaXYMax, trackDcaZMax, minSignImpXYSig); + } else { + jetProb = jettaggingutilities::getJetProbability(fSignImpXYSigIncJetMC, jet, tracks, trackDcaXYMax, trackDcaZMax, minSignImpXYSig); + } + } else { + if (origin == JetTaggingSpecies::charm) { + if (usepTcategorize) { + jetProb = jettaggingutilities::getJetProbability(vecfSignImpXYSigCharmJetMcCCDB, jet, tracks, trackDcaXYMax, trackDcaZMax, minSignImpXYSig); + } else { + jetProb = jettaggingutilities::getJetProbability(fSignImpXYSigCharmJetMC, jet, tracks, trackDcaXYMax, trackDcaZMax, minSignImpXYSig); + } + } else if (origin == JetTaggingSpecies::beauty) { + if (usepTcategorize) { + jetProb = jettaggingutilities::getJetProbability(vecfSignImpXYSigBeautyJetMcCCDB, jet, tracks, trackDcaXYMax, trackDcaZMax, minSignImpXYSig); + } else { + jetProb = jettaggingutilities::getJetProbability(fSignImpXYSigBeautyJetMC, jet, tracks, trackDcaXYMax, trackDcaZMax, minSignImpXYSig); + } + } else { + if (usepTcategorize) { + jetProb = jettaggingutilities::getJetProbability(vecfSignImpXYSigLfJetMcCCDB, jet, tracks, trackDcaXYMax, trackDcaZMax, minSignImpXYSig); + } else { + jetProb = jettaggingutilities::getJetProbability(fSignImpXYSigLfJetMC, jet, tracks, trackDcaXYMax, trackDcaZMax, minSignImpXYSig); + } + } + } + } + return jetProb; + } + + template + void evaluateTrackProbQA(int origin, T const& jet, U const& /*tracks*/, bool const& isMC = true) + { + for (const auto& track : jet.template tracks_as()) { + if (!jettaggingutilities::trackAcceptanceWithDca(track, trackDcaXYMax, trackDcaZMax)) + continue; + auto geoSign = jettaggingutilities::getGeoSign(jet, track); + float probTrack = -1; + if (!isMC) { + probTrack = jettaggingutilities::getTrackProbability(fSignImpXYSigData, track, minSignImpXYSig); + if (geoSign > 0) + registry.fill(HIST("h_pos_track_probability"), probTrack); + else + registry.fill(HIST("h_neg_track_probability"), probTrack); + } else { + if (useResoFuncFromIncJet) { + probTrack = jettaggingutilities::getTrackProbability(fSignImpXYSigIncJetMC, track, minSignImpXYSig); + } else { + if (origin == JetTaggingSpecies::charm) { + probTrack = jettaggingutilities::getTrackProbability(fSignImpXYSigCharmJetMC, track, minSignImpXYSig); + } + if (origin == JetTaggingSpecies::beauty) { + probTrack = jettaggingutilities::getTrackProbability(fSignImpXYSigBeautyJetMC, track, minSignImpXYSig); + } + if (origin == JetTaggingSpecies::lightflavour) { + probTrack = jettaggingutilities::getTrackProbability(fSignImpXYSigLfJetMC, track, minSignImpXYSig); + } + } + if (geoSign > 0) + registry.fill(HIST("h2_pos_track_probability_flavour"), probTrack, origin); + else + registry.fill(HIST("h2_neg_track_probability_flavour"), probTrack, origin); + } + } + } + + template + void fillTables(T const& jet, U const& tracks) + { + int origin = -1; + float jetProb = -1.0; + if constexpr (isMC) { + origin = jet.origin(); + } + if (useJetProb) { + if constexpr (isMC) { + jetProb = calculateJetProbability(origin, jet, tracks, isMC); + if (trackProbQA) { + evaluateTrackProbQA(origin, jet, tracks, isMC); + } + } else { + jetProb = calculateJetProbability(0, jet, tracks, isMC); + if (trackProbQA) { + evaluateTrackProbQA(0, jet, tracks, isMC); + } + } + } + if (doprocessAlgorithmGNN || doprocessAlgorithmGNNFull || doprocessAlgorithmGNNwExtra) { + if (jet.pt() >= jetpTMin) { + float dbRange; + if (scoreML[jet.globalIndex()] < dbMin) { + dbRange = 0.5; // underflow + } else if (scoreML[jet.globalIndex()] < dbMax) { + dbRange = 1.5; // in range + } else { + dbRange = 2.5; // overflow + } + registry.fill(HIST("h2_count_db"), 3.5, dbRange); // incl jet + if constexpr (isMC) { + switch (origin) { + case 2: + registry.fill(HIST("h_db_b"), scoreML[jet.globalIndex()]); + registry.fill(HIST("h2_count_db"), 0.5, dbRange); // b-jet + break; + case 1: + registry.fill(HIST("h_db_c"), scoreML[jet.globalIndex()]); + registry.fill(HIST("h2_count_db"), 1.5, dbRange); // c-jet + break; + case 0: + case 3: + registry.fill(HIST("h_db_lf"), scoreML[jet.globalIndex()]); + registry.fill(HIST("h2_count_db"), 2.5, dbRange); // lf-jet + break; + default: + LOGF(debug, "doprocessAlgorithmGNN, Unexpected origin value: %d (%d)", origin, jet.globalIndex()); + } + } + registry.fill(HIST("h2_pt_db"), jet.pt(), scoreML[jet.globalIndex()]); + } + // gnnTrackTable(scoreMLAux1[]) + // gnnVertexTable(scoreMLAux2[]) + } + taggingTable(decisionNonML[jet.globalIndex()], jetProb, scoreML[jet.globalIndex()]); + } + + HistogramRegistry registry{"registry", {}, OutputObjHandlingPolicy::AnalysisObject}; + void init(InitContext const&) + { + std::vector vecParamsData; + std::vector vecParamsIncJetMC; + std::vector vecParamsCharmJetMC; + std::vector vecParamsBeautyJetMC; + std::vector vecParamsLfJetMC; + + std::vector resoFuncDataCCDB; + std::vector resoFuncIncCCDB; + std::vector resoFuncBeautyCCDB; + std::vector resoFuncCharmCCDB; + std::vector resoFuncLfCCDB; + + ccdbApi.init(ccdbUrl); + + std::map metadata; + resoFuncMatch = resoFuncMatching; + + const int mIPmethodResolutionFunctionSize = 7; + + auto loadCCDBforIP = [&](const std::vector& paths, std::vector& targetVec, const std::string& name) { + if (paths.size() != mIPmethodResolutionFunctionSize) { + usepTcategorize.value = false; + LOG(info) << name << " does not have 7 entries. Disabling pT categorization (usepTcategorize = false)."; + resoFuncMatch = 0; + return; + } + for (int i = 0; i < mIPmethodResolutionFunctionSize; i++) { + targetVec.push_back(ccdbApi.retrieveFromTFileAny(paths[i], metadata, timestampCCDBForIP)); + } + }; + + if (usepTcategorize) { + switch (resoFuncMatch) { + case 6: + loadCCDBforIP(pathsCCDBforIPIncparameter, resoFuncIncCCDB, "pathsCCDBforIPIncparameter"); + break; + + case 7: + loadCCDBforIP(pathsCCDBforIPBeautyparameter, resoFuncBeautyCCDB, "pathsCCDBforIPBeautyparameter"); + loadCCDBforIP(pathsCCDBforIPCharmparameter, resoFuncCharmCCDB, "pathsCCDBforIPCharmparameter"); + loadCCDBforIP(pathsCCDBforIPLfparameter, resoFuncLfCCDB, "pathsCCDBforIPLfparameter"); + break; + + case 8: + loadCCDBforIP(pathsCCDBforIPDataparameter, resoFuncDataCCDB, "pathsCCDBforIPDataparameter"); + break; + + default: + LOG(info) << "resoFuncMatching is neither 6 nor 7, although usepTcategorize is set to true. Resetting resoFuncMatching to 0."; + resoFuncMatch = 0; + break; + } + } + + maxOrder = numCount + 1; // 0: untagged, >1 : N ordering + const int mIPmethodNumOfParameters = 9; + + // Set up the resolution function + switch (resoFuncMatch) { + case 0: + vecParamsData = (std::vector)paramsResoFuncData; + vecParamsCharmJetMC = (std::vector)paramsResoFuncCharmJetMC; + vecParamsBeautyJetMC = (std::vector)paramsResoFuncBeautyJetMC; + vecParamsLfJetMC = (std::vector)paramsResoFuncLfJetMC; + LOG(info) << "defined parameters of resolution function: custom"; + break; + case 1: + vecParamsData = (std::vector)paramsResoFuncData; + vecParamsIncJetMC = (std::vector)paramsResoFuncIncJetMC; + useResoFuncFromIncJet = true; + LOG(info) << "defined parameters of resolution function: custom & use inclusive distribution"; + break; + case 2: // TODO + vecParamsData = (std::vector)paramsResoFuncData; + vecParamsCharmJetMC = {282119.753, -0.065, 0.893, 11.608, 0.945, 8.029, 0.131, 6.244, 0.027}; + vecParamsBeautyJetMC = {74901.583, -0.082, 0.874, 10.332, 0.941, 7.352, 0.097, 6.220, 0.022}; + vecParamsLfJetMC = {1539435.343, -0.061, 0.896, 13.272, 1.034, 5.884, 0.004, 7.843, 0.090}; + LOG(info) << "defined parameters of resolution function: PYTHIA8, MB, LHC23d1k"; + break; + case 3: // TODO + vecParamsData = (std::vector)paramsResoFuncData; + vecParamsIncJetMC = {1908803.027, -0.059, 0.895, 13.467, 1.005, 8.867, 0.098, 6.929, 0.011}; + LOG(info) << "defined parameters of resolution function: PYTHIA8, MB, LHC23d1k & use inclusive distribution"; + useResoFuncFromIncJet = true; + break; + case 4: // TODO + vecParamsData = (std::vector)paramsResoFuncData; + vecParamsCharmJetMC = {743719.121, -0.960, -0.240, 13.765, 1.314, 10.761, 0.293, 8.538, 0.052}; + vecParamsBeautyJetMC = {88888.418, 0.256, 1.003, 10.185, 0.740, 8.216, 0.147, 7.228, 0.040}; + vecParamsLfJetMC = {414860.372, -1.000, 0.285, 14.561, 1.464, 11.693, 0.339, 9.183, 0.052}; + LOG(info) << "defined parameters of resolution function: PYTHIA8, JJ, weighted, LHC24g4"; + break; + case 5: // TODO + vecParamsData = (std::vector)paramsResoFuncData; + vecParamsIncJetMC = {2211391.862, 0.360, 1.028, 13.019, 0.650, 11.151, 0.215, 9.462, 0.044}; + LOG(info) << "defined parameters of resolution function: PYTHIA8, JJ, weighted, LHC24g4 & use inclusive distribution"; + useResoFuncFromIncJet = true; + break; + case 6: // TODO + vecParamsData = (std::vector)paramsResoFuncData; + for (size_t j = 0; j < resoFuncIncCCDB.size(); j++) { + std::vector params; + if (resoFuncIncCCDB[j]) { + for (int i = 0; i < mIPmethodNumOfParameters; i++) { + params.emplace_back(resoFuncIncCCDB[j]->GetParameter(i)); + } + } + vecParamsIncJetMcCCDB.emplace_back(params); + } + LOG(info) << "defined parameters of resolution function from CCDB"; + useResoFuncFromIncJet = true; + break; + case 7: // TODO + vecParamsData = (std::vector)paramsResoFuncData; + for (size_t j = 0; j < resoFuncBeautyCCDB.size(); j++) { + std::vector params; + if (resoFuncBeautyCCDB[j]) { + for (int i = 0; i < mIPmethodNumOfParameters; i++) { + params.emplace_back(resoFuncBeautyCCDB[j]->GetParameter(i)); + } + } + vecParamsBeautyJetMcCCDB.emplace_back(params); + } + for (size_t j = 0; j < resoFuncCharmCCDB.size(); j++) { + std::vector params; + if (resoFuncCharmCCDB[j]) { + for (int i = 0; i < mIPmethodNumOfParameters; i++) { + params.emplace_back(resoFuncCharmCCDB[j]->GetParameter(i)); + } + } + vecParamsCharmJetMcCCDB.emplace_back(params); + } + for (size_t j = 0; j < resoFuncLfCCDB.size(); j++) { + std::vector params; + if (resoFuncLfCCDB[j]) { + for (int i = 0; i < mIPmethodNumOfParameters; i++) { + params.emplace_back(resoFuncLfCCDB[j]->GetParameter(i)); + } + } + vecParamsLfJetMcCCDB.emplace_back(params); + } + LOG(info) << "defined parameters of resolution function from CCDB for each flavour"; + break; + case 8: + for (size_t j = 0; j < resoFuncDataCCDB.size(); j++) { + std::vector params; + if (resoFuncDataCCDB[j]) { + for (int i = 0; i < mIPmethodNumOfParameters; i++) { + params.emplace_back(resoFuncDataCCDB[j]->GetParameter(i)); + } + } + vecParamsDataJetCCDB.emplace_back(params); + } + LOG(info) << "defined parameters of resolution function from CCDB for data"; + break; + default: + LOG(fatal) << "undefined parameters of resolution function. Fix it!"; + break; + } + + fSignImpXYSigData = jettaggingutilities::setResolutionFunction(vecParamsData); + fSignImpXYSigIncJetMC = jettaggingutilities::setResolutionFunction(vecParamsIncJetMC); + fSignImpXYSigCharmJetMC = jettaggingutilities::setResolutionFunction(vecParamsCharmJetMC); + fSignImpXYSigBeautyJetMC = jettaggingutilities::setResolutionFunction(vecParamsBeautyJetMC); + fSignImpXYSigLfJetMC = jettaggingutilities::setResolutionFunction(vecParamsLfJetMC); + + for (const auto& params : vecParamsDataJetCCDB) { + vecfSignImpXYSigDataJetCCDB.emplace_back(jettaggingutilities::setResolutionFunction(params)); + } + for (const auto& params : vecParamsIncJetMcCCDB) { + vecfSignImpXYSigIncJetMcCCDB.emplace_back(jettaggingutilities::setResolutionFunction(params)); + } + for (const auto& params : vecParamsBeautyJetMcCCDB) { + vecfSignImpXYSigBeautyJetMcCCDB.emplace_back(jettaggingutilities::setResolutionFunction(params)); + } + for (const auto& params : vecParamsCharmJetMcCCDB) { + vecfSignImpXYSigCharmJetMcCCDB.emplace_back(jettaggingutilities::setResolutionFunction(params)); + } + for (const auto& params : vecParamsLfJetMcCCDB) { + vecfSignImpXYSigLfJetMcCCDB.emplace_back(jettaggingutilities::setResolutionFunction(params)); + } + + // Use QA for effectivness of track probability + if (trackProbQA) { + AxisSpec trackProbabilityAxis = {binTrackProbability, "Track proability"}; + AxisSpec jetFlavourAxis = {binJetFlavour, "Jet flavour"}; + if (doprocessFillTables) { + if (isMCD) { + registry.add("h2_pos_track_probability_flavour", "positive track probability", {HistType::kTH2F, {{trackProbabilityAxis}, {jetFlavourAxis}}}); + registry.add("h2_neg_track_probability_flavour", "negative track probability", {HistType::kTH2F, {{trackProbabilityAxis}, {jetFlavourAxis}}}); + } else { + registry.add("h_pos_track_probability", "positive track probability", {HistType::kTH1F, {{trackProbabilityAxis}}}); + registry.add("h_neg_track_probability", "negative track probability", {HistType::kTH1F, {{trackProbabilityAxis}}}); + } + } + } + + if (doprocessAlgorithmML || doprocessAlgorithmGNN || doprocessAlgorithmGNNFull || doprocessAlgorithmGNNwExtra) { + bMlResponse.configure(binsPtMl, cutsMl, cutDirMl, nClassesMl); + if (loadModelsFromCCDB) { + ccdbApi.init(ccdbUrl); + bMlResponse.setModelPathsCCDB(onnxFileNames, ccdbApi, modelPathsCCDB, timestampCCDB); + } else { + bMlResponse.setModelPathsLocal(onnxFileNames); + } + bMlResponse.cacheInputFeaturesIndices(namesInputFeatures); + bMlResponse.init(); + } + + if (doprocessAlgorithmGNN || doprocessAlgorithmGNNFull || doprocessAlgorithmGNNwExtra) { + tensorAlloc = o2::analysis::GNNBjetAllocator(nJetFeat.value, nTrkFeat.value, nClassesMl.value, nTrkOrigin.value, transformFeatureJetMean.value, transformFeatureJetStdev.value, transformFeatureTrkMean.value, transformFeatureTrkStdev.value, nJetConst, tfFuncTypeGNN.value); + + registry.add("h2_count_db", "#it{D}_{b} underflow/overflow;Jet flavour;#it{D}_{b} range", {HistType::kTH2F, {{4, 0., 4.}, {3, 0., 3.}}}); + auto h2CountDb = registry.get(HIST("h2_count_db")); + h2CountDb->GetXaxis()->SetBinLabel(1, "b-jet"); + h2CountDb->GetXaxis()->SetBinLabel(2, "c-jet"); + h2CountDb->GetXaxis()->SetBinLabel(3, "lf-jet"); + h2CountDb->GetXaxis()->SetBinLabel(4, "incl jet"); + h2CountDb->GetYaxis()->SetBinLabel(1, "underflow"); + h2CountDb->GetYaxis()->SetBinLabel(2, "in range"); + h2CountDb->GetYaxis()->SetBinLabel(3, "overflow"); + + registry.add("h_db_b", "#it{D}_{b} b-jet;#it{D}_{b}", {HistType::kTH1F, {{50, -20., 30.}}}); + registry.add("h_db_c", "#it{D}_{b} c-jet;#it{D}_{b}", {HistType::kTH1F, {{50, -20., 30.}}}); + registry.add("h_db_lf", "#it{D}_{b} lf-jet;#it{D}_{b}", {HistType::kTH1F, {{50, -20., 30.}}}); + registry.add("h2_pt_db", "#it{p}_{T} vs. #it{D}_{b};#it{p}_{T}^{ch jet} (GeV/#it{c}^{2});#it{D}_{b}", {HistType::kTH2F, {{100, 0., 200.}, {50, -20., 30.}}}); + } + + if (doprocessAlgorithmGNNFull) { + } + } + + template + void analyzeJetAlgorithmML(AnyJets const& alljets, AnyTracks const& allTracks, SecondaryVertices const& allSVs) + { + for (const auto& analysisJet : alljets) { + + std::vector tracksParams; + std::vector svsParams; + + jettaggingutilities::analyzeJetSVInfo4ML(analysisJet, allTracks, allSVs, svsParams, svPtMin, svReductionFactor); + jettaggingutilities::analyzeJetTrackInfo4ML(analysisJet, allTracks, allSVs, tracksParams, trackPtMin, trackDcaXYMax, trackDcaZMax); + + int nSVs = analysisJet.template secondaryVertices_as().size(); + + jettaggingutilities::BJetParams jetparam = {analysisJet.pt(), analysisJet.eta(), analysisJet.phi(), static_cast(tracksParams.size()), static_cast(nSVs), analysisJet.mass()}; + tracksParams.resize(nJetConst); // resize to the number of inputs of the ML + svsParams.resize(nJetConst); // resize to the number of inputs of the ML + + std::vector output; + + if (bMlResponse.getInputShape().size() > 1) { + auto inputML = bMlResponse.getInputFeatures2D(jetparam, tracksParams, svsParams); + bMlResponse.isSelectedMl(inputML, analysisJet.pt(), output); + } else { + auto inputML = bMlResponse.getInputFeatures1D(jetparam, tracksParams, svsParams); + bMlResponse.isSelectedMl(inputML, analysisJet.pt(), output); + } + + if (bMlResponse.getOutputNodes() > 1) { + auto mDb = [](const std::vector& scores, float fC) { + return std::log(scores[2] / (fC * scores[1] + (1 - fC) * scores[0])); + }; + + scoreML[analysisJet.globalIndex()] = useDb ? mDb(output, fC) : output[2]; // 2 is the b-jet index + } else { + scoreML[analysisJet.globalIndex()] = output[0]; + } + } + } + + template + void analyzeJetAlgorithmMLnoSV(AnyJets const& alljets, AnyTracks const& allTracks) + { + for (const auto& analysisJet : alljets) { + + std::vector tracksParams; + std::vector svsParams; + + jettaggingutilities::analyzeJetTrackInfo4MLnoSV(analysisJet, allTracks, tracksParams, trackPtMin, trackDcaXYMax, trackDcaZMax); + + jettaggingutilities::BJetParams jetparam = {analysisJet.pt(), analysisJet.eta(), analysisJet.phi(), static_cast(tracksParams.size()), 0, analysisJet.mass()}; + tracksParams.resize(nJetConst); // resize to the number of inputs of the ML + + std::vector output; + + if (bMlResponse.getInputShape().size() > 1) { + auto inputML = bMlResponse.getInputFeatures2D(jetparam, tracksParams, svsParams); + bMlResponse.isSelectedMl(inputML, analysisJet.pt(), output); + } else { + auto inputML = bMlResponse.getInputFeatures1D(jetparam, tracksParams, svsParams); + bMlResponse.isSelectedMl(inputML, analysisJet.pt(), output); + } + + scoreML[analysisJet.globalIndex()] = output[0]; + } + } + + template + void analyzeJetAlgorithmGNNwExtra(AnyJets const& jets, AnyTracks const& tracks, AnyOriginalTracks const& origTracks) + { + for (const auto& jet : jets) { + std::vector> trkFeat; + jettaggingutilities::analyzeJetTrackInfo4GNNwExtra(jet, tracks, origTracks, trkFeat, trackPtMin, trackDcaXYMax, trackDcaZMax, nJetConst); + + std::vector jetFeat{jet.pt(), jet.phi(), jet.eta(), jet.mass()}; + + if (trkFeat.size() > 0) { + std::vector feat; + std::vector gnnInput; + tensorAlloc.getGNNInput(jetFeat, trkFeat, feat, gnnInput); + + auto modelOutput = bMlResponse.getModelOutput(gnnInput, 0); + float db = jettaggingutilities::getDb(modelOutput, fC); + if (!std::isnan(db)) { + scoreML[jet.globalIndex()] = db; + } else { + scoreML[jet.globalIndex()] = 999.; + LOGF(debug, "doprocessAlgorithmGNNwExtra, Db is NaN (%d)", jet.globalIndex()); + } + } else { + scoreML[jet.globalIndex()] = -999.; + LOGF(debug, "doprocessAlgorithmGNNwExtra, trkFeat.size() <= 0 (%d)", jet.globalIndex()); + } + } + } + + template + void analyzeJetAlgorithmGNN(AnyJets const& jets, AnyTracks const& tracks) + { + for (const auto& jet : jets) { + std::vector> trkFeat; + jettaggingutilities::analyzeJetTrackInfo4GNN(jet, tracks, trkFeat, trackPtMin, trackDcaXYMax, trackDcaZMax, nJetConst); + + std::vector jetFeat{jet.pt(), jet.phi(), jet.eta(), jet.mass()}; + + if (trkFeat.size() > 0) { + std::vector feat; + std::vector gnnInput; + tensorAlloc.getGNNInput(jetFeat, trkFeat, feat, gnnInput); + + auto modelOutput = bMlResponse.getModelOutput(gnnInput, 0); + 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()); + } + } + } + + template + void analyzeJetAlgorithmGNNFull(AnyJets const& jets, AnyTracks const& tracks) + { + for (const auto& jet : jets) { + std::vector> trkFeat; + jettaggingutilities::analyzeJetTrackInfo4GNN(jet, tracks, trkFeat, trackPtMin, trackDcaXYMax, trackDcaZMax, nJetConst); + + std::vector jetFeat{jet.pt(), jet.phi(), jet.eta(), jet.mass()}; + + if (trkFeat.size() > 0) { + std::vector feat; + std::vector gnnInput; + tensorAlloc.getGNNInput(jetFeat, trkFeat, feat, gnnInput); + + auto modelOutputs = bMlResponse.getModelOutputs(gnnInput, 0); + + float db = jettaggingutilities::getDb(modelOutputs[idxGnnJet], fC); + if (!std::isnan(db)) { + scoreML[jet.globalIndex()] = db; + } else { + scoreML[jet.globalIndex()] = 999.; + LOGF(debug, "doprocessAlgorithmGNN, Db is NaN (%d)", jet.globalIndex()); + } + + // modelOutputs[idxGnnTrack], modelOutputs[idxGnnVertex] + } else { + scoreML[jet.globalIndex()] = -999.; + LOGF(debug, "doprocessAlgorithmGNN, trkFeat.size() <= 0 (%d)", jet.globalIndex()); + } + } + } + + void processDummy(aod::JetCollisions const&) + { + } + PROCESS_SWITCH(JetTaggerHFTask, processDummy, "Dummy process", true); + + void processSetup(JetTable const& jets) + { + decisionNonML.clear(); + scoreML.clear(); + decisionNonML.resize(jets.size()); + scoreML.resize(jets.size()); + } + PROCESS_SWITCH(JetTaggerHFTask, processSetup, "Setup initialization and size of jets for filling table", false); + + void processIP(JetTable const& jets, JetTracksExt const& tracks) + { + for (const auto& jet : jets) { + uint16_t bit = jettaggingutilities::setTaggingIPBit(jet, tracks, trackDcaXYMax, trackDcaZMax, tagPointForIP); + decisionNonML[jet.globalIndex()] |= bit; + } + } + PROCESS_SWITCH(JetTaggerHFTask, processIP, "Fill tagging decision for jet with the IP algorithm", false); + + void processSV(soa::Join const& jets, SVTable const& prongs) + { + for (const auto& jet : jets) { + uint16_t bit = jettaggingutilities::setTaggingSVBit(jet, prongs, prongChi2PCAMin, prongChi2PCAMax, prongsigmaLxyMax, prongIPxyMin, prongIPxyMax, svDispersionMax, tagPointForSV); + decisionNonML[jet.globalIndex()] |= bit; + } + } + PROCESS_SWITCH(JetTaggerHFTask, processSV, "Fill tagging decision for jets with the SV", false); + + void processAlgorithmML(soa::Join const& allJets, JetTracksExt const& allTracks, SVTable const& allSVs) + { + analyzeJetAlgorithmML(allJets, allTracks, allSVs); + } + PROCESS_SWITCH(JetTaggerHFTask, processAlgorithmML, "Fill ML evaluation score for charged jets", false); + + void processAlgorithmMLnoSV(JetTable const& allJets, JetTracksExt const& allTracks) + { + analyzeJetAlgorithmMLnoSV(allJets, allTracks); + } + PROCESS_SWITCH(JetTaggerHFTask, processAlgorithmMLnoSV, "Fill ML evaluation score for charged jets but without using SVs", false); + + void processAlgorithmGNNwExtra(JetTable const& jets, JetTracksExt const& jtracks, OriginalTracks const& origTracks) + { + analyzeJetAlgorithmGNNwExtra(jets, jtracks, origTracks); + } + PROCESS_SWITCH(JetTaggerHFTask, processAlgorithmGNNwExtra, "Fill GNN evaluation score (D_b) for charged jets with extra input features", false); + + void processAlgorithmGNN(JetTable const& jets, JetTracksExt const& jtracks) + { + analyzeJetAlgorithmGNN(jets, jtracks); + } + PROCESS_SWITCH(JetTaggerHFTask, processAlgorithmGNN, "Fill GNN evaluation score (D_b) for charged jets", false); + + void processAlgorithmGNNFull(JetTable const& jets, JetTracksExt const& jtracks) + { + analyzeJetAlgorithmGNNFull(jets, jtracks); + } + PROCESS_SWITCH(JetTaggerHFTask, processAlgorithmGNNFull, "Fill GNN evaluation score (D_b) for charged jets, + GNN track/vertex predictions", false); + + void processFillTables(std::conditional_t, JetTable>::iterator const& jet, JetTracksExt const& tracks) + { + fillTables(jet, tracks); + } + PROCESS_SWITCH(JetTaggerHFTask, processFillTables, "Fill Tables for tagging decision, jet probability, and ML score on charged jets", false); +}; + +using JetTaggerhfDataCharged = JetTaggerHFTask, aod::DataSecondaryVertexNProngIndices, aod::DataSecondaryVertexNProngs, aod::ChargedJetTags>; +using JetTaggerhfMcdCharged = JetTaggerHFTask, aod::MCDSecondaryVertexNProngIndices, aod::MCDSecondaryVertexNProngs, aod::ChargedMCDetectorLevelJetTags>; +using JetTaggerhfDataChargedGNNFull = JetTaggerHFTask, aod::DataSecondaryVertexNProngIndices, aod::DataSecondaryVertexNProngs, aod::ChargedJetTags>; +using JetTaggerhfMcdChargedGNNFull = JetTaggerHFTask, aod::MCDSecondaryVertexNProngIndices, aod::MCDSecondaryVertexNProngs, aod::ChargedMCDetectorLevelJetTags>; + +WorkflowSpec defineDataProcessing(ConfigContext const& cfgc) +{ + + return WorkflowSpec{ + adaptAnalysisTask( + cfgc, SetDefaultProcesses{}, TaskName{"jet-taggerhf-data-charged"}), // o2-linter: disable=name/o2-task (custom device name for templated task alias) + adaptAnalysisTask( + cfgc, SetDefaultProcesses{}, TaskName{"jet-taggerhf-mcd-charged"})}; // o2-linter: disable=name/o2-task (custom device name for templated task alias) +} diff --git a/Tools/ML/MlResponse.h b/Tools/ML/MlResponse.h index 413534f48cb..863753b1d46 100644 --- a/Tools/ML/MlResponse.h +++ b/Tools/ML/MlResponse.h @@ -241,6 +241,27 @@ class MlResponse return output; } + /// Get vector of vectors with predictions from every output head of a multi-output model (e.g. GNN) + /// \param input a vector containing the values of features used in the model + /// \param nModel is the model index + /// \return one vector per output head, in the order of the ONNX model's output nodes + template + std::vector> getModelOutputs(T1& input, const T2& nModel) + { + if (nModel < 0 || static_cast(nModel) >= mModels.size()) { + LOG(fatal) << "Model index " << nModel << " is out of range! The number of initialised models is " << mModels.size() << ". Please check your configurables."; + } + + const int numInputNodes = mModels[nModel].getNumInputNodes(); + const int numInputFeatures = static_cast(input.size()); + + if (numInputNodes != numInputFeatures && numInputNodes >= 0) { + LOG(fatal) << "Number of input nodes in the model " << mPaths[nModel] << " is different from the number of input features to be tested (" << numInputNodes << " vs " << numInputFeatures << ")"; + } + + return mModels[nModel].template evalModelFull(input); + } + /// ML selections /// \param input is the input features /// \param candVar is the variable value (e.g. pT) used to select which model to use diff --git a/Tools/ML/model.h b/Tools/ML/model.h index d94491f8c38..00d3d2d552a 100644 --- a/Tools/ML/model.h +++ b/Tools/ML/model.h @@ -144,6 +144,79 @@ class OnnxModel evalModel(inputTensors, output); } + // Evaluate model and retrieve the FULL set of output tensors (e.g. multi-head models such as GNNs). + // Unlike evalModel(), which returns only the last output tensor as a raw pointer, this + // copies every output tensor's data into an owned vector, sized from that tensor's own shape. + template + std::vector> evalModelFull(std::vector& input) + { + std::vector> results; + try { + const Ort::RunOptions runOptions; + std::vector inputNamesChar(mInputNames.size(), nullptr); + std::transform(std::begin(mInputNames), std::end(mInputNames), std::begin(inputNamesChar), + [&](const std::string& str) { return str.c_str(); }); + + std::vector outputNamesChar(mOutputNames.size(), nullptr); + std::transform(std::begin(mOutputNames), std::end(mOutputNames), std::begin(outputNamesChar), + [&](const std::string& str) { return str.c_str(); }); + + auto outputTensors = mSession->Run(runOptions, inputNamesChar.data(), input.data(), input.size(), outputNamesChar.data(), outputNamesChar.size()); + if (outputTensors.size() != mOutputNames.size()) { + LOG(fatal) << "Number of output tensors: " << outputTensors.size() << " does not agree with the model specified size: " << mOutputNames.size(); + } + + results.reserve(outputTensors.size()); + for (auto& tensor : outputTensors) { + const std::size_t nElements = tensor.GetTensorTypeAndShapeInfo().GetElementCount(); + const T* data = tensor.GetTensorData(); + results.emplace_back(data, data + nElements); + } + } catch (const Ort::Exception& exception) { + LOG(error) << "Error running model inference: " << exception.what(); + } + return results; + } + + // Flat-vector input overload, mirrors evalModel(std::vector&) + template + std::vector> evalModelFull(std::vector& input) + { + const int64_t size = input.size(); + assert(size % mInputShapes[0][1] == 0); + std::vector inputShape{size / mInputShapes[0][1], mInputShapes[0][1]}; + std::vector inputTensors; + Ort::MemoryInfo memInfo = Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault); + inputTensors.emplace_back(Ort::Value::CreateTensor(memInfo, input.data(), size, inputShape.data(), inputShape.size())); + return evalModelFull(inputTensors); + } + + // 2D-vector input overload, mirrors evalModel(std::vector>&) + template + std::vector> evalModelFull(std::vector>& input) + { + std::vector inputTensors; + Ort::MemoryInfo memInfo = Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault); + + for (std::size_t iinput = 0; iinput < input.size(); iinput++) { + [[maybe_unused]] int totalSize = 1; + int64_t size = input[iinput].size(); + for (std::size_t idim = 1; idim < mInputShapes[iinput].size(); idim++) { + totalSize *= mInputShapes[iinput][idim]; + } + assert(size % totalSize == 0); + + std::vector inputShape{static_cast(size / totalSize)}; + for (std::size_t idim = 1; idim < mInputShapes[iinput].size(); idim++) { + inputShape.push_back(mInputShapes[iinput][idim]); + } + + inputTensors.emplace_back(Ort::Value::CreateTensor(memInfo, input[iinput].data(), size, inputShape.data(), inputShape.size())); + } + + return evalModelFull(inputTensors); + } + // Reset session void resetSession() { From 6a23a4b6b15c5606c173e2b1630d6b57e230dc7d Mon Sep 17 00:00:00 2001 From: Changhwan Choi Date: Wed, 9 Sep 2026 18:26:38 +0900 Subject: [PATCH 2/7] jetGnnInference.cxx to utilise full outputs of GNN b-jet tagger --- PWGJE/TableProducer/CMakeLists.txt | 5 + PWGJE/TableProducer/jetGnnInference.cxx | 573 ++---------------------- 2 files changed, 30 insertions(+), 548 deletions(-) diff --git a/PWGJE/TableProducer/CMakeLists.txt b/PWGJE/TableProducer/CMakeLists.txt index c789aff9d49..8f4c523d23f 100644 --- a/PWGJE/TableProducer/CMakeLists.txt +++ b/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-inference + SOURCES jetGnnInference.cxx + PUBLIC_LINK_LIBRARIES O2::Framework O2Physics::PWGJECore O2Physics::AnalysisCore O2Physics::MLCore + COMPONENT_NAME Analysis) + o2physics_add_dpl_workflow(estimator-rho SOURCES rhoEstimator.cxx REUSE_FROM JetDataModelPCH diff --git a/PWGJE/TableProducer/jetGnnInference.cxx b/PWGJE/TableProducer/jetGnnInference.cxx index 3249c6cb96b..50709fa371a 100644 --- a/PWGJE/TableProducer/jetGnnInference.cxx +++ b/PWGJE/TableProducer/jetGnnInference.cxx @@ -56,62 +56,31 @@ using namespace o2; using namespace o2::framework; using namespace o2::framework::expressions; -template +template struct JetGnnInferenceTask { static constexpr double DefaultCutsMl[1][2] = {{0.5, 0.5}}; Produces taggingTable; - Produces gnnTrackTable; - Produces gnnVertexTable; + if constexpr (!std::is_same_v) { + Produces gnnTrackTable; + } + if constexpr (!std::is_same_v) { + Produces gnnVertexTable; + } // configuration topological cut for track and sv Configurable trackDcaXYMax{"trackDcaXYMax", 1, "minimum DCA xy acceptance for tracks [cm]"}; Configurable trackDcaZMax{"trackDcaZMax", 2, "minimum DCA z acceptance for tracks [cm]"}; - // Configurable prongsigmaLxyMax{"prongsigmaLxyMax", 100, "maximum sigma of decay length of prongs on xy plane"}; - // Configurable prongsigmaLxyzMax{"prongsigmaLxyzMax", 100, "maximum sigma of decay length of prongs on xyz plane"}; - // Configurable prongIPxyMin{"prongIPxyMin", 0.008, "maximum impact paramter of prongs on xy plane [cm]"}; - // Configurable prongIPxyMax{"prongIPxyMax", 1, "minimum impact parmeter of prongs on xy plane [cm]"}; - // Configurable prongChi2PCAMin{"prongChi2PCAMin", 4, "minimum Chi2 PCA of decay length of prongs"}; - // Configurable prongChi2PCAMax{"prongChi2PCAMax", 100, "maximum Chi2 PCA of decay length of prongs"}; - // Configurable svDispersionMax{"svDispersionMax", 1, "maximum dispersion of sv"}; - - // // configuration about IP method - // Configurable useJetProb{"useJetProb", false, "fill table for track counting algorithm"}; - // Configurable trackProbQA{"trackProbQA", false, "fill track probability histograms separately for geometric positive and negative tracks for QA"}; - // Configurable numCount{"numCount", 3, "number of track counting"}; - // Configurable resoFuncMatching{"resoFuncMatching", 0, "matching parameters of resolution function as MC samble (0: custom, 1: custom & inc, 2: MB, 3: MB & inc, 4: JJ, 5: JJ & inc)"}; - // Configurable> pathsCCDBforIPDataparameter{"pathsCCDBforIPDataparameter", std::vector{"Users/l/leehy/LHC24g4/f_inclusive_0"}, "Paths for fitting parameters of resolution functions of data for IP method on CCDB"}; - // Configurable> pathsCCDBforIPIncparameter{"pathsCCDBforIPIncparameter", std::vector{"Users/l/leehy/LHC24g4/f_inclusive_0"}, "Paths for fitting parameters of resolution functions of inclusive for IP method on CCDB"}; - // Configurable> pathsCCDBforIPBeautyparameter{"pathsCCDBforIPBeautyparameter", std::vector{"Users/l/leehy/LHC24g4/f_inclusive_0"}, "Paths for fitting parameters of resolution functions of beauty for IP method on CCDB"}; - // Configurable> pathsCCDBforIPCharmparameter{"pathsCCDBforIPCharmparameter", std::vector{"Users/l/leehy/LHC24g4/f_inclusive_0"}, "Paths for fitting parameters of resolution functions of charm for IP method on CCDB"}; - // Configurable> pathsCCDBforIPLfparameter{"pathsCCDBforIPLfparameter", std::vector{"Users/l/leehy/LHC24g4/f_inclusive_0"}, "Paths for fitting parameters of resolution functions of light flavour for IP method on CCDB"}; - // Configurable usepTcategorize{"usepTcategorize", false, "p_T categorize TF1 function with Inclusive jet"}; - // Configurable> paramsResoFuncData{"paramsResoFuncData", std::vector{-1.0}, "parameters of gaus(0)+expo(3)+expo(5)+expo(7))"}; - // Configurable> paramsResoFuncIncJetMC{"paramsResoFuncIncJetMC", std::vector{-1.0}, "parameters of gaus(0)+expo(3)+expo(5)+expo(7)))"}; - // Configurable> paramsResoFuncCharmJetMC{"paramsResoFuncCharmJetMC", std::vector{-1.0}, "parameters of gaus(0)+expo(3)+expo(5)+expo(7)))"}; - // Configurable> paramsResoFuncBeautyJetMC{"paramsResoFuncBeautyJetMC", std::vector{-1.0}, "parameters of gaus(0)+expo(3)+expo(5)+expo(7)))"}; - // Configurable> paramsResoFuncLfJetMC{"paramsResoFuncLfJetMC", std::vector{-1.0}, "parameters of gaus(0)+expo(3)+expo(5)+expo(7)))"}; - // Configurable minSignImpXYSig{"minSignImpXYSig", -40.0, "minimum of signed impact parameter significance"}; - // Configurable tagPointForIP{"tagPointForIP", 2.5, "tagging working point for IP"}; - // Configurable tagPointForIPxyz{"tagPointForIPxyz", 2.5, "tagging working point for IP xyz"}; - // Configurable timestampCCDBForIP{"timestampCCDBForIP", -1, "timestamp of the resolution function file for IP method used to query in CCDB"}; - // // configuration about SV method - // Configurable tagPointForSV{"tagPointForSV", 40, "tagging working point for SV"}; - // Configurable tagPointForSVxyz{"tagPointForSVxyz", 40, "tagging working point for SV xyz"}; // ML configuration Configurable nJetConst{"nJetConst", 100, "maximum number of jet consistuents to be used for ML evaluation"}; Configurable trackPtMin{"trackPtMin", 0.15, "minimum track pT"}; - // Configurable svPtMin{"svPtMin", 0.5, "minimum SV pT"}; - - // Configurable svReductionFactor{"svReductionFactor", 1.0, "factor for how many SVs to keep"}; Configurable> binsPtMl{"binsPtMl", std::vector{5., 1000.}, "pT bin limits for ML application"}; Configurable> cutDirMl{"cutDirMl", std::vector{cuts_ml::CutSmaller, cuts_ml::CutNot}, "Whether to reject score values greater or smaller than the threshold"}; Configurable> cutsMl{"cutsMl", {DefaultCutsMl[0], 1, 2, {"pT bin 0"}, {"score for default b-jet tagging", "uncer 1"}}, "ML selections per pT bin"}; Configurable nClassesMl{"nClassesMl", 3, "Number of classes in ML model"}; Configurable> namesInputFeatures{"namesInputFeatures", std::vector{"feature1", "feature2"}, "Names of ML model input features"}; - // Configurable useDb{"useDb", false, "Flag to use DB for ML model instead of the score"}; Configurable ccdbUrl{"ccdbUrl", "http://alice-ccdb.cern.ch", "url of the ccdb repository"}; Configurable> modelPathsCCDB{"modelPathsCCDB", std::vector{"Users/c/chchoi"}, "Paths of models on CCDB"}; @@ -121,8 +90,8 @@ struct JetGnnInferenceTask { // GNN configuration Configurable jetpTMin{"jetpTMin", 5., "minimum jet pT"}; - Configurable dbMin{"dbMin", -10., "minimum GNN Db"}; - Configurable dbMax{"dbMax", 20., "maximum GNN Db"}; + Configurable dbMin{"dbMin", -15., "minimum GNN Db"}; + Configurable dbMax{"dbMax", 15., "maximum GNN Db"}; Configurable fC{"fC", 0.018, "Parameter f_c for D_b calculation"}; Configurable nJetFeat{"nJetFeat", 4, "Number of jet GNN input features"}; Configurable nTrkFeat{"nTrkFeat", 13, "Number of track GNN input features"}; @@ -139,11 +108,7 @@ struct JetGnnInferenceTask { Configurable> transformFeatureTrkStdev{"transformFeatureTrkStdev", std::vector{-999}, "Stdev values for each GNN input feature (track)"}; - Configurable tfFuncTypeGNN{"tfFuncTypeGNN", "linear", "Transformation function type for GNN"}; - - // axis spec - ConfigurableAxis binTrackProbability{"binTrackProbability", {100, 0.f, 1.f}, ""}; - ConfigurableAxis binJetFlavour{"binJetFlavour", {6, -0.5, 5.5}, ""}; + Configurable tfFuncTypeGNN{"tfFuncTypeGNN", "asinh", "Transformation function type for GNN"}; o2::analysis::MlResponseHfTagging bMlResponse; o2::ccdb::CcdbApi ccdbApi; @@ -151,111 +116,10 @@ struct JetGnnInferenceTask { using JetTracksExt = soa::Join; using OriginalTracks = soa::Join; - bool useResoFuncFromIncJet = false; - int maxOrder = -1; - int resoFuncMatch = 0; - - std::unique_ptr fSignImpXYSigData = nullptr; - std::unique_ptr fSignImpXYSigIncJetMC = nullptr; - std::unique_ptr fSignImpXYSigCharmJetMC = nullptr; - std::unique_ptr fSignImpXYSigBeautyJetMC = nullptr; - std::unique_ptr fSignImpXYSigLfJetMC = nullptr; - - std::vector> vecParamsDataJetCCDB; - std::vector> vecParamsIncJetMcCCDB; - std::vector> vecParamsBeautyJetMcCCDB; - std::vector> vecParamsCharmJetMcCCDB; - std::vector> vecParamsLfJetMcCCDB; - - std::vector> vecfSignImpXYSigDataJetCCDB; - std::vector> vecfSignImpXYSigIncJetMcCCDB; - std::vector> vecfSignImpXYSigCharmJetMcCCDB; - std::vector> vecfSignImpXYSigBeautyJetMcCCDB; - std::vector> vecfSignImpXYSigLfJetMcCCDB; - - std::vector decisionNonML; std::vector scoreML; o2::analysis::GNNBjetAllocator tensorAlloc; - template - float calculateJetProbability(int origin, T const& jet, U const& tracks, bool const& isMC = false) - { - float jetProb = -1.0; - if (!isMC) { - if (usepTcategorize) { - jetProb = jettaggingutilities::getJetProbability(vecfSignImpXYSigDataJetCCDB, jet, tracks, trackDcaXYMax, trackDcaZMax, minSignImpXYSig); - } else { - jetProb = jettaggingutilities::getJetProbability(fSignImpXYSigData, jet, tracks, trackDcaXYMax, trackDcaZMax, minSignImpXYSig); - } - } else { - if (useResoFuncFromIncJet) { - if (usepTcategorize) { - jetProb = jettaggingutilities::getJetProbability(vecfSignImpXYSigIncJetMcCCDB, jet, tracks, trackDcaXYMax, trackDcaZMax, minSignImpXYSig); - } else { - jetProb = jettaggingutilities::getJetProbability(fSignImpXYSigIncJetMC, jet, tracks, trackDcaXYMax, trackDcaZMax, minSignImpXYSig); - } - } else { - if (origin == JetTaggingSpecies::charm) { - if (usepTcategorize) { - jetProb = jettaggingutilities::getJetProbability(vecfSignImpXYSigCharmJetMcCCDB, jet, tracks, trackDcaXYMax, trackDcaZMax, minSignImpXYSig); - } else { - jetProb = jettaggingutilities::getJetProbability(fSignImpXYSigCharmJetMC, jet, tracks, trackDcaXYMax, trackDcaZMax, minSignImpXYSig); - } - } else if (origin == JetTaggingSpecies::beauty) { - if (usepTcategorize) { - jetProb = jettaggingutilities::getJetProbability(vecfSignImpXYSigBeautyJetMcCCDB, jet, tracks, trackDcaXYMax, trackDcaZMax, minSignImpXYSig); - } else { - jetProb = jettaggingutilities::getJetProbability(fSignImpXYSigBeautyJetMC, jet, tracks, trackDcaXYMax, trackDcaZMax, minSignImpXYSig); - } - } else { - if (usepTcategorize) { - jetProb = jettaggingutilities::getJetProbability(vecfSignImpXYSigLfJetMcCCDB, jet, tracks, trackDcaXYMax, trackDcaZMax, minSignImpXYSig); - } else { - jetProb = jettaggingutilities::getJetProbability(fSignImpXYSigLfJetMC, jet, tracks, trackDcaXYMax, trackDcaZMax, minSignImpXYSig); - } - } - } - } - return jetProb; - } - - template - void evaluateTrackProbQA(int origin, T const& jet, U const& /*tracks*/, bool const& isMC = true) - { - for (const auto& track : jet.template tracks_as()) { - if (!jettaggingutilities::trackAcceptanceWithDca(track, trackDcaXYMax, trackDcaZMax)) - continue; - auto geoSign = jettaggingutilities::getGeoSign(jet, track); - float probTrack = -1; - if (!isMC) { - probTrack = jettaggingutilities::getTrackProbability(fSignImpXYSigData, track, minSignImpXYSig); - if (geoSign > 0) - registry.fill(HIST("h_pos_track_probability"), probTrack); - else - registry.fill(HIST("h_neg_track_probability"), probTrack); - } else { - if (useResoFuncFromIncJet) { - probTrack = jettaggingutilities::getTrackProbability(fSignImpXYSigIncJetMC, track, minSignImpXYSig); - } else { - if (origin == JetTaggingSpecies::charm) { - probTrack = jettaggingutilities::getTrackProbability(fSignImpXYSigCharmJetMC, track, minSignImpXYSig); - } - if (origin == JetTaggingSpecies::beauty) { - probTrack = jettaggingutilities::getTrackProbability(fSignImpXYSigBeautyJetMC, track, minSignImpXYSig); - } - if (origin == JetTaggingSpecies::lightflavour) { - probTrack = jettaggingutilities::getTrackProbability(fSignImpXYSigLfJetMC, track, minSignImpXYSig); - } - } - if (geoSign > 0) - registry.fill(HIST("h2_pos_track_probability_flavour"), probTrack, origin); - else - registry.fill(HIST("h2_neg_track_probability_flavour"), probTrack, origin); - } - } - } - template void fillTables(T const& jet, U const& tracks) { @@ -264,20 +128,7 @@ struct JetGnnInferenceTask { if constexpr (isMC) { origin = jet.origin(); } - if (useJetProb) { - if constexpr (isMC) { - jetProb = calculateJetProbability(origin, jet, tracks, isMC); - if (trackProbQA) { - evaluateTrackProbQA(origin, jet, tracks, isMC); - } - } else { - jetProb = calculateJetProbability(0, jet, tracks, isMC); - if (trackProbQA) { - evaluateTrackProbQA(0, jet, tracks, isMC); - } - } - } - if (doprocessAlgorithmGNN || doprocessAlgorithmGNNFull || doprocessAlgorithmGNNwExtra) { + if (doprocessAlgorithmGNN) { if (jet.pt() >= jetpTMin) { float dbRange; if (scoreML[jet.globalIndex()] < dbMin) { @@ -312,210 +163,13 @@ struct JetGnnInferenceTask { // gnnTrackTable(scoreMLAux1[]) // gnnVertexTable(scoreMLAux2[]) } - taggingTable(decisionNonML[jet.globalIndex()], jetProb, scoreML[jet.globalIndex()]); + taggingTable(0, jetProb, scoreML[jet.globalIndex()]); } HistogramRegistry registry{"registry", {}, OutputObjHandlingPolicy::AnalysisObject}; void init(InitContext const&) { - std::vector vecParamsData; - std::vector vecParamsIncJetMC; - std::vector vecParamsCharmJetMC; - std::vector vecParamsBeautyJetMC; - std::vector vecParamsLfJetMC; - - std::vector resoFuncDataCCDB; - std::vector resoFuncIncCCDB; - std::vector resoFuncBeautyCCDB; - std::vector resoFuncCharmCCDB; - std::vector resoFuncLfCCDB; - - ccdbApi.init(ccdbUrl); - - std::map metadata; - resoFuncMatch = resoFuncMatching; - - const int mIPmethodResolutionFunctionSize = 7; - - auto loadCCDBforIP = [&](const std::vector& paths, std::vector& targetVec, const std::string& name) { - if (paths.size() != mIPmethodResolutionFunctionSize) { - usepTcategorize.value = false; - LOG(info) << name << " does not have 7 entries. Disabling pT categorization (usepTcategorize = false)."; - resoFuncMatch = 0; - return; - } - for (int i = 0; i < mIPmethodResolutionFunctionSize; i++) { - targetVec.push_back(ccdbApi.retrieveFromTFileAny(paths[i], metadata, timestampCCDBForIP)); - } - }; - - if (usepTcategorize) { - switch (resoFuncMatch) { - case 6: - loadCCDBforIP(pathsCCDBforIPIncparameter, resoFuncIncCCDB, "pathsCCDBforIPIncparameter"); - break; - - case 7: - loadCCDBforIP(pathsCCDBforIPBeautyparameter, resoFuncBeautyCCDB, "pathsCCDBforIPBeautyparameter"); - loadCCDBforIP(pathsCCDBforIPCharmparameter, resoFuncCharmCCDB, "pathsCCDBforIPCharmparameter"); - loadCCDBforIP(pathsCCDBforIPLfparameter, resoFuncLfCCDB, "pathsCCDBforIPLfparameter"); - break; - - case 8: - loadCCDBforIP(pathsCCDBforIPDataparameter, resoFuncDataCCDB, "pathsCCDBforIPDataparameter"); - break; - - default: - LOG(info) << "resoFuncMatching is neither 6 nor 7, although usepTcategorize is set to true. Resetting resoFuncMatching to 0."; - resoFuncMatch = 0; - break; - } - } - - maxOrder = numCount + 1; // 0: untagged, >1 : N ordering - const int mIPmethodNumOfParameters = 9; - - // Set up the resolution function - switch (resoFuncMatch) { - case 0: - vecParamsData = (std::vector)paramsResoFuncData; - vecParamsCharmJetMC = (std::vector)paramsResoFuncCharmJetMC; - vecParamsBeautyJetMC = (std::vector)paramsResoFuncBeautyJetMC; - vecParamsLfJetMC = (std::vector)paramsResoFuncLfJetMC; - LOG(info) << "defined parameters of resolution function: custom"; - break; - case 1: - vecParamsData = (std::vector)paramsResoFuncData; - vecParamsIncJetMC = (std::vector)paramsResoFuncIncJetMC; - useResoFuncFromIncJet = true; - LOG(info) << "defined parameters of resolution function: custom & use inclusive distribution"; - break; - case 2: // TODO - vecParamsData = (std::vector)paramsResoFuncData; - vecParamsCharmJetMC = {282119.753, -0.065, 0.893, 11.608, 0.945, 8.029, 0.131, 6.244, 0.027}; - vecParamsBeautyJetMC = {74901.583, -0.082, 0.874, 10.332, 0.941, 7.352, 0.097, 6.220, 0.022}; - vecParamsLfJetMC = {1539435.343, -0.061, 0.896, 13.272, 1.034, 5.884, 0.004, 7.843, 0.090}; - LOG(info) << "defined parameters of resolution function: PYTHIA8, MB, LHC23d1k"; - break; - case 3: // TODO - vecParamsData = (std::vector)paramsResoFuncData; - vecParamsIncJetMC = {1908803.027, -0.059, 0.895, 13.467, 1.005, 8.867, 0.098, 6.929, 0.011}; - LOG(info) << "defined parameters of resolution function: PYTHIA8, MB, LHC23d1k & use inclusive distribution"; - useResoFuncFromIncJet = true; - break; - case 4: // TODO - vecParamsData = (std::vector)paramsResoFuncData; - vecParamsCharmJetMC = {743719.121, -0.960, -0.240, 13.765, 1.314, 10.761, 0.293, 8.538, 0.052}; - vecParamsBeautyJetMC = {88888.418, 0.256, 1.003, 10.185, 0.740, 8.216, 0.147, 7.228, 0.040}; - vecParamsLfJetMC = {414860.372, -1.000, 0.285, 14.561, 1.464, 11.693, 0.339, 9.183, 0.052}; - LOG(info) << "defined parameters of resolution function: PYTHIA8, JJ, weighted, LHC24g4"; - break; - case 5: // TODO - vecParamsData = (std::vector)paramsResoFuncData; - vecParamsIncJetMC = {2211391.862, 0.360, 1.028, 13.019, 0.650, 11.151, 0.215, 9.462, 0.044}; - LOG(info) << "defined parameters of resolution function: PYTHIA8, JJ, weighted, LHC24g4 & use inclusive distribution"; - useResoFuncFromIncJet = true; - break; - case 6: // TODO - vecParamsData = (std::vector)paramsResoFuncData; - for (size_t j = 0; j < resoFuncIncCCDB.size(); j++) { - std::vector params; - if (resoFuncIncCCDB[j]) { - for (int i = 0; i < mIPmethodNumOfParameters; i++) { - params.emplace_back(resoFuncIncCCDB[j]->GetParameter(i)); - } - } - vecParamsIncJetMcCCDB.emplace_back(params); - } - LOG(info) << "defined parameters of resolution function from CCDB"; - useResoFuncFromIncJet = true; - break; - case 7: // TODO - vecParamsData = (std::vector)paramsResoFuncData; - for (size_t j = 0; j < resoFuncBeautyCCDB.size(); j++) { - std::vector params; - if (resoFuncBeautyCCDB[j]) { - for (int i = 0; i < mIPmethodNumOfParameters; i++) { - params.emplace_back(resoFuncBeautyCCDB[j]->GetParameter(i)); - } - } - vecParamsBeautyJetMcCCDB.emplace_back(params); - } - for (size_t j = 0; j < resoFuncCharmCCDB.size(); j++) { - std::vector params; - if (resoFuncCharmCCDB[j]) { - for (int i = 0; i < mIPmethodNumOfParameters; i++) { - params.emplace_back(resoFuncCharmCCDB[j]->GetParameter(i)); - } - } - vecParamsCharmJetMcCCDB.emplace_back(params); - } - for (size_t j = 0; j < resoFuncLfCCDB.size(); j++) { - std::vector params; - if (resoFuncLfCCDB[j]) { - for (int i = 0; i < mIPmethodNumOfParameters; i++) { - params.emplace_back(resoFuncLfCCDB[j]->GetParameter(i)); - } - } - vecParamsLfJetMcCCDB.emplace_back(params); - } - LOG(info) << "defined parameters of resolution function from CCDB for each flavour"; - break; - case 8: - for (size_t j = 0; j < resoFuncDataCCDB.size(); j++) { - std::vector params; - if (resoFuncDataCCDB[j]) { - for (int i = 0; i < mIPmethodNumOfParameters; i++) { - params.emplace_back(resoFuncDataCCDB[j]->GetParameter(i)); - } - } - vecParamsDataJetCCDB.emplace_back(params); - } - LOG(info) << "defined parameters of resolution function from CCDB for data"; - break; - default: - LOG(fatal) << "undefined parameters of resolution function. Fix it!"; - break; - } - - fSignImpXYSigData = jettaggingutilities::setResolutionFunction(vecParamsData); - fSignImpXYSigIncJetMC = jettaggingutilities::setResolutionFunction(vecParamsIncJetMC); - fSignImpXYSigCharmJetMC = jettaggingutilities::setResolutionFunction(vecParamsCharmJetMC); - fSignImpXYSigBeautyJetMC = jettaggingutilities::setResolutionFunction(vecParamsBeautyJetMC); - fSignImpXYSigLfJetMC = jettaggingutilities::setResolutionFunction(vecParamsLfJetMC); - - for (const auto& params : vecParamsDataJetCCDB) { - vecfSignImpXYSigDataJetCCDB.emplace_back(jettaggingutilities::setResolutionFunction(params)); - } - for (const auto& params : vecParamsIncJetMcCCDB) { - vecfSignImpXYSigIncJetMcCCDB.emplace_back(jettaggingutilities::setResolutionFunction(params)); - } - for (const auto& params : vecParamsBeautyJetMcCCDB) { - vecfSignImpXYSigBeautyJetMcCCDB.emplace_back(jettaggingutilities::setResolutionFunction(params)); - } - for (const auto& params : vecParamsCharmJetMcCCDB) { - vecfSignImpXYSigCharmJetMcCCDB.emplace_back(jettaggingutilities::setResolutionFunction(params)); - } - for (const auto& params : vecParamsLfJetMcCCDB) { - vecfSignImpXYSigLfJetMcCCDB.emplace_back(jettaggingutilities::setResolutionFunction(params)); - } - - // Use QA for effectivness of track probability - if (trackProbQA) { - AxisSpec trackProbabilityAxis = {binTrackProbability, "Track proability"}; - AxisSpec jetFlavourAxis = {binJetFlavour, "Jet flavour"}; - if (doprocessFillTables) { - if (isMCD) { - registry.add("h2_pos_track_probability_flavour", "positive track probability", {HistType::kTH2F, {{trackProbabilityAxis}, {jetFlavourAxis}}}); - registry.add("h2_neg_track_probability_flavour", "negative track probability", {HistType::kTH2F, {{trackProbabilityAxis}, {jetFlavourAxis}}}); - } else { - registry.add("h_pos_track_probability", "positive track probability", {HistType::kTH1F, {{trackProbabilityAxis}}}); - registry.add("h_neg_track_probability", "negative track probability", {HistType::kTH1F, {{trackProbabilityAxis}}}); - } - } - } - - if (doprocessAlgorithmML || doprocessAlgorithmGNN || doprocessAlgorithmGNNFull || doprocessAlgorithmGNNwExtra) { + if (doprocessAlgorithmGNN) { bMlResponse.configure(binsPtMl, cutsMl, cutDirMl, nClassesMl); if (loadModelsFromCCDB) { ccdbApi.init(ccdbUrl); @@ -525,9 +179,7 @@ struct JetGnnInferenceTask { } bMlResponse.cacheInputFeaturesIndices(namesInputFeatures); bMlResponse.init(); - } - if (doprocessAlgorithmGNN || doprocessAlgorithmGNNFull || doprocessAlgorithmGNNwExtra) { tensorAlloc = o2::analysis::GNNBjetAllocator(nJetFeat.value, nTrkFeat.value, nClassesMl.value, nTrkOrigin.value, transformFeatureJetMean.value, transformFeatureJetStdev.value, transformFeatureTrkMean.value, transformFeatureTrkStdev.value, nJetConst, tfFuncTypeGNN.value); registry.add("h2_count_db", "#it{D}_{b} underflow/overflow;Jet flavour;#it{D}_{b} range", {HistType::kTH2F, {{4, 0., 4.}, {3, 0., 3.}}}); @@ -545,104 +197,6 @@ struct JetGnnInferenceTask { registry.add("h_db_lf", "#it{D}_{b} lf-jet;#it{D}_{b}", {HistType::kTH1F, {{50, -20., 30.}}}); registry.add("h2_pt_db", "#it{p}_{T} vs. #it{D}_{b};#it{p}_{T}^{ch jet} (GeV/#it{c}^{2});#it{D}_{b}", {HistType::kTH2F, {{100, 0., 200.}, {50, -20., 30.}}}); } - - if (doprocessAlgorithmGNNFull) { - } - } - - template - void analyzeJetAlgorithmML(AnyJets const& alljets, AnyTracks const& allTracks, SecondaryVertices const& allSVs) - { - for (const auto& analysisJet : alljets) { - - std::vector tracksParams; - std::vector svsParams; - - jettaggingutilities::analyzeJetSVInfo4ML(analysisJet, allTracks, allSVs, svsParams, svPtMin, svReductionFactor); - jettaggingutilities::analyzeJetTrackInfo4ML(analysisJet, allTracks, allSVs, tracksParams, trackPtMin, trackDcaXYMax, trackDcaZMax); - - int nSVs = analysisJet.template secondaryVertices_as().size(); - - jettaggingutilities::BJetParams jetparam = {analysisJet.pt(), analysisJet.eta(), analysisJet.phi(), static_cast(tracksParams.size()), static_cast(nSVs), analysisJet.mass()}; - tracksParams.resize(nJetConst); // resize to the number of inputs of the ML - svsParams.resize(nJetConst); // resize to the number of inputs of the ML - - std::vector output; - - if (bMlResponse.getInputShape().size() > 1) { - auto inputML = bMlResponse.getInputFeatures2D(jetparam, tracksParams, svsParams); - bMlResponse.isSelectedMl(inputML, analysisJet.pt(), output); - } else { - auto inputML = bMlResponse.getInputFeatures1D(jetparam, tracksParams, svsParams); - bMlResponse.isSelectedMl(inputML, analysisJet.pt(), output); - } - - if (bMlResponse.getOutputNodes() > 1) { - auto mDb = [](const std::vector& scores, float fC) { - return std::log(scores[2] / (fC * scores[1] + (1 - fC) * scores[0])); - }; - - scoreML[analysisJet.globalIndex()] = useDb ? mDb(output, fC) : output[2]; // 2 is the b-jet index - } else { - scoreML[analysisJet.globalIndex()] = output[0]; - } - } - } - - template - void analyzeJetAlgorithmMLnoSV(AnyJets const& alljets, AnyTracks const& allTracks) - { - for (const auto& analysisJet : alljets) { - - std::vector tracksParams; - std::vector svsParams; - - jettaggingutilities::analyzeJetTrackInfo4MLnoSV(analysisJet, allTracks, tracksParams, trackPtMin, trackDcaXYMax, trackDcaZMax); - - jettaggingutilities::BJetParams jetparam = {analysisJet.pt(), analysisJet.eta(), analysisJet.phi(), static_cast(tracksParams.size()), 0, analysisJet.mass()}; - tracksParams.resize(nJetConst); // resize to the number of inputs of the ML - - std::vector output; - - if (bMlResponse.getInputShape().size() > 1) { - auto inputML = bMlResponse.getInputFeatures2D(jetparam, tracksParams, svsParams); - bMlResponse.isSelectedMl(inputML, analysisJet.pt(), output); - } else { - auto inputML = bMlResponse.getInputFeatures1D(jetparam, tracksParams, svsParams); - bMlResponse.isSelectedMl(inputML, analysisJet.pt(), output); - } - - scoreML[analysisJet.globalIndex()] = output[0]; - } - } - - template - void analyzeJetAlgorithmGNNwExtra(AnyJets const& jets, AnyTracks const& tracks, AnyOriginalTracks const& origTracks) - { - for (const auto& jet : jets) { - std::vector> trkFeat; - jettaggingutilities::analyzeJetTrackInfo4GNNwExtra(jet, tracks, origTracks, trkFeat, trackPtMin, trackDcaXYMax, trackDcaZMax, nJetConst); - - std::vector jetFeat{jet.pt(), jet.phi(), jet.eta(), jet.mass()}; - - if (trkFeat.size() > 0) { - std::vector feat; - std::vector gnnInput; - tensorAlloc.getGNNInput(jetFeat, trkFeat, feat, gnnInput); - - auto modelOutput = bMlResponse.getModelOutput(gnnInput, 0); - float db = jettaggingutilities::getDb(modelOutput, fC); - if (!std::isnan(db)) { - scoreML[jet.globalIndex()] = db; - } else { - scoreML[jet.globalIndex()] = 999.; - LOGF(debug, "doprocessAlgorithmGNNwExtra, Db is NaN (%d)", jet.globalIndex()); - } - } else { - scoreML[jet.globalIndex()] = -999.; - LOGF(debug, "doprocessAlgorithmGNNwExtra, trkFeat.size() <= 0 (%d)", jet.globalIndex()); - } - } } template @@ -674,118 +228,41 @@ struct JetGnnInferenceTask { } } - template - void analyzeJetAlgorithmGNNFull(AnyJets const& jets, AnyTracks const& tracks) - { - for (const auto& jet : jets) { - std::vector> trkFeat; - jettaggingutilities::analyzeJetTrackInfo4GNN(jet, tracks, trkFeat, trackPtMin, trackDcaXYMax, trackDcaZMax, nJetConst); - - std::vector jetFeat{jet.pt(), jet.phi(), jet.eta(), jet.mass()}; - - if (trkFeat.size() > 0) { - std::vector feat; - std::vector gnnInput; - tensorAlloc.getGNNInput(jetFeat, trkFeat, feat, gnnInput); - - auto modelOutputs = bMlResponse.getModelOutputs(gnnInput, 0); - - float db = jettaggingutilities::getDb(modelOutputs[idxGnnJet], fC); - if (!std::isnan(db)) { - scoreML[jet.globalIndex()] = db; - } else { - scoreML[jet.globalIndex()] = 999.; - LOGF(debug, "doprocessAlgorithmGNN, Db is NaN (%d)", jet.globalIndex()); - } - - // modelOutputs[idxGnnTrack], modelOutputs[idxGnnVertex] - } else { - scoreML[jet.globalIndex()] = -999.; - LOGF(debug, "doprocessAlgorithmGNN, trkFeat.size() <= 0 (%d)", jet.globalIndex()); - } - } - } - void processDummy(aod::JetCollisions const&) { } - PROCESS_SWITCH(JetTaggerHFTask, processDummy, "Dummy process", true); + PROCESS_SWITCH(JetGnnInferenceTask, processDummy, "Dummy process", true); void processSetup(JetTable const& jets) { - decisionNonML.clear(); scoreML.clear(); - decisionNonML.resize(jets.size()); scoreML.resize(jets.size()); } - PROCESS_SWITCH(JetTaggerHFTask, processSetup, "Setup initialization and size of jets for filling table", false); - - void processIP(JetTable const& jets, JetTracksExt const& tracks) - { - for (const auto& jet : jets) { - uint16_t bit = jettaggingutilities::setTaggingIPBit(jet, tracks, trackDcaXYMax, trackDcaZMax, tagPointForIP); - decisionNonML[jet.globalIndex()] |= bit; - } - } - PROCESS_SWITCH(JetTaggerHFTask, processIP, "Fill tagging decision for jet with the IP algorithm", false); - - void processSV(soa::Join const& jets, SVTable const& prongs) - { - for (const auto& jet : jets) { - uint16_t bit = jettaggingutilities::setTaggingSVBit(jet, prongs, prongChi2PCAMin, prongChi2PCAMax, prongsigmaLxyMax, prongIPxyMin, prongIPxyMax, svDispersionMax, tagPointForSV); - decisionNonML[jet.globalIndex()] |= bit; - } - } - PROCESS_SWITCH(JetTaggerHFTask, processSV, "Fill tagging decision for jets with the SV", false); - - void processAlgorithmML(soa::Join const& allJets, JetTracksExt const& allTracks, SVTable const& allSVs) - { - analyzeJetAlgorithmML(allJets, allTracks, allSVs); - } - PROCESS_SWITCH(JetTaggerHFTask, processAlgorithmML, "Fill ML evaluation score for charged jets", false); - - void processAlgorithmMLnoSV(JetTable const& allJets, JetTracksExt const& allTracks) - { - analyzeJetAlgorithmMLnoSV(allJets, allTracks); - } - PROCESS_SWITCH(JetTaggerHFTask, processAlgorithmMLnoSV, "Fill ML evaluation score for charged jets but without using SVs", false); - - void processAlgorithmGNNwExtra(JetTable const& jets, JetTracksExt const& jtracks, OriginalTracks const& origTracks) - { - analyzeJetAlgorithmGNNwExtra(jets, jtracks, origTracks); - } - PROCESS_SWITCH(JetTaggerHFTask, processAlgorithmGNNwExtra, "Fill GNN evaluation score (D_b) for charged jets with extra input features", false); + PROCESS_SWITCH(JetGnnInferenceTask, processSetup, "Setup initialization and size of jets for filling table", false); void processAlgorithmGNN(JetTable const& jets, JetTracksExt const& jtracks) { analyzeJetAlgorithmGNN(jets, jtracks); } - PROCESS_SWITCH(JetTaggerHFTask, processAlgorithmGNN, "Fill GNN evaluation score (D_b) for charged jets", false); - - void processAlgorithmGNNFull(JetTable const& jets, JetTracksExt const& jtracks) - { - analyzeJetAlgorithmGNNFull(jets, jtracks); - } - PROCESS_SWITCH(JetTaggerHFTask, processAlgorithmGNNFull, "Fill GNN evaluation score (D_b) for charged jets, + GNN track/vertex predictions", false); + PROCESS_SWITCH(JetGnnInferenceTask, processAlgorithmGNN, "Fill GNN evaluation score (D_b) for charged jets", false); void processFillTables(std::conditional_t, JetTable>::iterator const& jet, JetTracksExt const& tracks) { fillTables(jet, tracks); } - PROCESS_SWITCH(JetTaggerHFTask, processFillTables, "Fill Tables for tagging decision, jet probability, and ML score on charged jets", false); + PROCESS_SWITCH(JetGnnInferenceTask, processFillTables, "Fill Tables for tagging decision, jet probability, and ML score on charged jets", false); }; -using JetTaggerhfDataCharged = JetTaggerHFTask, aod::DataSecondaryVertexNProngIndices, aod::DataSecondaryVertexNProngs, aod::ChargedJetTags>; -using JetTaggerhfMcdCharged = JetTaggerHFTask, aod::MCDSecondaryVertexNProngIndices, aod::MCDSecondaryVertexNProngs, aod::ChargedMCDetectorLevelJetTags>; -using JetTaggerhfDataChargedGNNFull = JetTaggerHFTask, aod::DataSecondaryVertexNProngIndices, aod::DataSecondaryVertexNProngs, aod::ChargedJetTags>; -using JetTaggerhfMcdChargedGNNFull = JetTaggerHFTask, aod::MCDSecondaryVertexNProngIndices, aod::MCDSecondaryVertexNProngs, aod::ChargedMCDetectorLevelJetTags>; +using JetGnnInferenceDataCharged = JetGnnInferenceTask, aod::ChargedJetTags>; +using JetGnnInferenceMcdCharged = JetGnnInferenceTask, aod::ChargedMCDetectorLevelJetTags>; +// using JetGnnInferenceDataChargedFull = JetGnnInferenceTask, aod::ChargedJetTags, ?, ?>; +// using JetGnnInferenceMcdChargedFull = JetGnnInferenceTask, aod::ChargedMCDetectorLevelJetTags, ?, ?>; WorkflowSpec defineDataProcessing(ConfigContext const& cfgc) { - return WorkflowSpec{ - adaptAnalysisTask( - cfgc, SetDefaultProcesses{}, TaskName{"jet-taggerhf-data-charged"}), // o2-linter: disable=name/o2-task (custom device name for templated task alias) - adaptAnalysisTask( - cfgc, SetDefaultProcesses{}, TaskName{"jet-taggerhf-mcd-charged"})}; // o2-linter: disable=name/o2-task (custom device name for templated task alias) + adaptAnalysisTask( + cfgc, SetDefaultProcesses{}, TaskName{"jet-gnn-inference-data-charged"}), // o2-linter: disable=name/o2-task (custom device name for templated task alias) + adaptAnalysisTask( + cfgc, SetDefaultProcesses{}, TaskName{"jet-gnn-inference-mcd-charged"})}; // o2-linter: disable=name/o2-task (custom device name for templated task alias) } From 10b0edc0173d657c6d0481f2d444b7ad15b26736 Mon Sep 17 00:00:00 2001 From: Changhwan Choi Date: Fri, 30 Oct 2026 00:06:28 +0900 Subject: [PATCH 3/7] pause --- PWGJE/DataModel/JetTagging.h | 98 +++++++++++++++++++++++++ PWGJE/TableProducer/jetGnnInference.cxx | 53 +++++++++---- 2 files changed, 137 insertions(+), 14 deletions(-) diff --git a/PWGJE/DataModel/JetTagging.h b/PWGJE/DataModel/JetTagging.h index 9ce0ddadb6f..033198f3484 100644 --- a/PWGJE/DataModel/JetTagging.h +++ b/PWGJE/DataModel/JetTagging.h @@ -30,6 +30,7 @@ #include #include #include +#include namespace o2::aod { @@ -165,6 +166,103 @@ DECLARE_SOA_DYNAMIC_COLUMN(IsTagged, isTagged, [](uint16_t bit, BJetTaggingMetho JETTAGGING_TABLES_DEF(Charged, "C"); 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 of the original track this in-jet track index corresponds to +DECLARE_SOA_COLUMN(TrackIndex, trackIndex, uint8_t); //! in-jet index i (0..N-1) of this constituent track, matching the order of jet.template tracks_as() +DECLARE_SOA_COLUMN(EdgeIndex, edgeIndex, uint8_t); //! unique in-jet edge index for the (i,j) constituent-track pair, i std::pair { + int i = 0; + int offset = 0; + while (offset + (n - 1 - i) <= static_cast(edgeIndex)) { + offset += (n - 1 - i); + i++; + } + int j = i + 1 + (static_cast(edgeIndex) - offset); + return std::pair{static_cast(i), static_cast(j)}; + }); +} // namespace jetgnnindex + +namespace jetgnnlabel +{ +DECLARE_SOA_COLUMN(Origin, origin, uint8_t); //! trkOrigin from vertexClustering() (0: not 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(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) +} // namespace jetgnnscore + +#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(_jet_type_##GNNTrack, tracks); \ + } \ + DECLARE_SOA_TABLE(_jet_type_##GNNTrackIndices, "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); \ + 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); + +// 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"); +DECLARE_JET_GNNTRACKINDEX_TABLE(ChargedMCDetectorLevelJet, "CDJGNT"); +DECLARE_JET_GNNEDGEINDEX_TABLE(ChargedJet, "CJGNE"); +DECLARE_JET_GNNEDGEINDEX_TABLE(ChargedMCDetectorLevelJet, "CDJGNE"); + +#undef DECLARE_JET_GNNTRACKINDEX_TABLE +#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). +DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetTrackOriginDef, "AOD", "CDJTOD", jetgnnlabel::Origin); +DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetVertexDef, "AOD", "CDJTVD", jetgnnlabel::VertexIndex); +DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetEdgeDef, "AOD", "CDJEDF", jetgnnlabel::SameVertex); + +// GNN ScoreML outputs, track (node) level, data and MC detector level. Thin tables: join positionally with +// ChargedJetGNNTrack / ChargedMCDetectorLevelJetGNNTrack. +DECLARE_SOA_TABLE(ChargedJetTrackTags, "AOD", "CJTRKT", jetgnnscore::ScoreML); +DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetTrackTags, "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); + +// Convenience joined types combining the backbone (in-jet ordering / original-track reference) with the thin +// content tables, so `.template tracks_as()` etc. can be used directly. +using ChargedJetGNNTrackLabels = o2::soa::Join; +using ChargedJetGNNEdgeLabels = o2::soa::Join; +using ChargedMCDetectorLevelJetGNNTrackLabels = o2::soa::Join; +using ChargedMCDetectorLevelJetGNNEdgeLabels = o2::soa::Join; + } // namespace o2::aod #endif // PWGJE_DATAMODEL_JETTAGGING_H_ diff --git a/PWGJE/TableProducer/jetGnnInference.cxx b/PWGJE/TableProducer/jetGnnInference.cxx index 50709fa371a..e63159336b1 100644 --- a/PWGJE/TableProducer/jetGnnInference.cxx +++ b/PWGJE/TableProducer/jetGnnInference.cxx @@ -56,17 +56,24 @@ using namespace o2; using namespace o2::framework; using namespace o2::framework::expressions; +namespace jet_gnn_inference +{ +// Helper struct to define optional produces for JetGnnInferenceTask +template +struct OptionalProduces : Produces { +}; +template <> +struct OptionalProduces { +}; +}; // namespace jet_gnn_inference + template struct JetGnnInferenceTask { static constexpr double DefaultCutsMl[1][2] = {{0.5, 0.5}}; Produces taggingTable; - if constexpr (!std::is_same_v) { - Produces gnnTrackTable; - } - if constexpr (!std::is_same_v) { - Produces gnnVertexTable; - } + jet_gnn_inference::OptionalProduces gnnTrackTable; + jet_gnn_inference::OptionalProduces gnnVertexTable; // configuration topological cut for track and sv Configurable trackDcaXYMax{"trackDcaXYMax", 1, "minimum DCA xy acceptance for tracks [cm]"}; @@ -167,6 +174,7 @@ struct JetGnnInferenceTask { } HistogramRegistry registry{"registry", {}, OutputObjHandlingPolicy::AnalysisObject}; + void init(InitContext const&) { if (doprocessAlgorithmGNN) { @@ -228,6 +236,12 @@ struct JetGnnInferenceTask { } } + template + void analyzeJetAlgorithmGNNBatch(AnyJets const& jets, AnyTracks const& tracks) + { + // Implement batch inference logic here + } + void processDummy(aod::JetCollisions const&) { } @@ -246,6 +260,13 @@ struct JetGnnInferenceTask { } PROCESS_SWITCH(JetGnnInferenceTask, processAlgorithmGNN, "Fill GNN evaluation score (D_b) for charged jets", false); + void processAlgorithmGNNBatch(aod::JetCollisions::iterator const& /*coll*/, JetTable const& jets, JetTracksExt const& jtracks) + { + // Batch inference of GNN model per collision + analyzeJetAlgorithmGNNBatch(jets, jtracks); + } + PROCESS_SWITCH(JetGnnInferenceTask, processAlgorithmGNNBatch, "Fill GNN evaluation score (D_b) for charged jets, batch inference", false); + void processFillTables(std::conditional_t, JetTable>::iterator const& jet, JetTracksExt const& tracks) { fillTables(jet, tracks); @@ -253,16 +274,20 @@ struct JetGnnInferenceTask { PROCESS_SWITCH(JetGnnInferenceTask, processFillTables, "Fill Tables for tagging decision, jet probability, and ML score on charged jets", false); }; -using JetGnnInferenceDataCharged = JetGnnInferenceTask, aod::ChargedJetTags>; -using JetGnnInferenceMcdCharged = JetGnnInferenceTask, aod::ChargedMCDetectorLevelJetTags>; -// using JetGnnInferenceDataChargedFull = JetGnnInferenceTask, aod::ChargedJetTags, ?, ?>; -// using JetGnnInferenceMcdChargedFull = JetGnnInferenceTask, aod::ChargedMCDetectorLevelJetTags, ?, ?>; +using JetGnnTagInferenceDataCharged = JetGnnInferenceTask, aod::ChargedJetTags>; +using JetGnnTagInferenceMcdCharged = JetGnnInferenceTask, aod::ChargedMCDetectorLevelJetTags>; +using JetGnnFullInferenceDataCharged = JetGnnInferenceTask, aod::ChargedJetTags, ChargedJetGNNTrack, ChargedJetTrackTags, ChargedJetGNNEdge, ChargedJetEdgeTags>; +using JetGnnFullInferenceMcdCharged = JetGnnInferenceTask, aod::ChargedMCDetectorLevelJetTags, ChargedMCDetectorLevelJetGNNTrack, ChargedMCDetectorLevelJetTrackTags, ChargedMCDetectorLevelJetGNNEdge, ChargedMCDetectorLevelJetEdgeTags>; WorkflowSpec defineDataProcessing(ConfigContext const& cfgc) { return WorkflowSpec{ - adaptAnalysisTask( - cfgc, SetDefaultProcesses{}, TaskName{"jet-gnn-inference-data-charged"}), // o2-linter: disable=name/o2-task (custom device name for templated task alias) - adaptAnalysisTask( - cfgc, SetDefaultProcesses{}, TaskName{"jet-gnn-inference-mcd-charged"})}; // o2-linter: disable=name/o2-task (custom device name for templated task alias) + adaptAnalysisTask( + cfgc, SetDefaultProcesses{}, TaskName{"jet-gnn-tag-inference-data-charged"}), // o2-linter: disable=name/o2-task (custom device name for templated task alias) + adaptAnalysisTask( + cfgc, SetDefaultProcesses{}, TaskName{"jet-gnn-tag-inference-mcd-charged"}), // o2-linter: disable=name/o2-task (custom device name for templated task alias) + adaptAnalysisTask( + cfgc, SetDefaultProcesses{}, TaskName{"jet-gnn-full-inference-data-charged"}), // o2-linter: disable=name/o2-task (custom device name for templated task alias) + adaptAnalysisTask( + cfgc, SetDefaultProcesses{}, TaskName{"jet-gnn-full-inference-mcd-charged"})}; // o2-linter: disable=name/o2-task (custom device name for templated task alias) } From 8e1810939102dcbbaed89de3a6c30bcc14ca1926 Mon Sep 17 00:00:00 2001 From: Changhwan Choi Date: Mon, 9 Nov 2026 11:31:45 +0900 Subject: [PATCH 4/7] searchUpToQuark for bjetTreeCreator --- PWGJE/DataModel/JetTagging.h | 2 +- PWGJE/Tasks/bjetTreeCreator.cxx | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/PWGJE/DataModel/JetTagging.h b/PWGJE/DataModel/JetTagging.h index 033198f3484..e687328f82b 100644 --- a/PWGJE/DataModel/JetTagging.h +++ b/PWGJE/DataModel/JetTagging.h @@ -187,7 +187,7 @@ DECLARE_SOA_DYNAMIC_COLUMN(TrackIds, trackIds, //! inverse of the triangular edg namespace jetgnnlabel { -DECLARE_SOA_COLUMN(Origin, origin, uint8_t); //! trkOrigin from vertexClustering() (0: not physical primary, 1: charm, 2: beauty, 3: primary vertex, 4: other secondary) +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(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 diff --git a/PWGJE/Tasks/bjetTreeCreator.cxx b/PWGJE/Tasks/bjetTreeCreator.cxx index 7ab7bc11bc6..186842c0153 100644 --- a/PWGJE/Tasks/bjetTreeCreator.cxx +++ b/PWGJE/Tasks/bjetTreeCreator.cxx @@ -279,6 +279,7 @@ struct BJetTreeCreator { Configurable produceTree{"produceTree", true, "produce the jet TTree"}; Configurable vtxRes{"vtxRes", 0.01, "Vertex position resolution (cluster size) for GNN vertex predictions (cm)"}; + Configurable searchUpToQuark{"searchUpToQuark", true, "Flag whether to search for track origins up to the quark level"}; Configurable trainingDatasetRatioParam{"trainingDatasetRatioParam", 0, "Parameter for splitting training/evaluation datasets by collisionId"}; Configurable produceExtrbTableGnn{"produceExtrbTableGnn", false, "Flag whether to produce the extra table with track quality variables for GNN training"}; @@ -846,7 +847,7 @@ struct BJetTreeCreator { //+ TrackLabelMap trkLabels{{"trkVtxIndex", {}}, {"trkOrigin", {}}}; - int nVertices = jettaggingutilities::vertexClustering(collision.template mcCollision_as(), analysisJet, allTracks, MCParticles, origParticles, trkLabels, true, vtxRes, trackPtMin); + int nVertices = jettaggingutilities::vertexClustering(collision.template mcCollision_as(), analysisJet, allTracks, MCParticles, origParticles, trkLabels, searchUpToQuark, vtxRes, trackPtMin); analyzeJetTrackInfoForGNN(collision, analysisJet, allTracks, origTracks, indicesTracks, jetFlavor, eventWeight, &trkLabels); registry.fill(HIST("h2_jetMass_jetpT"), analysisJet.pt(), analysisJet.mass(), eventWeight); From ac87630e2fc8e634bd0b61014ecc7d8ff616045c Mon Sep 17 00:00:00 2001 From: Changhwan Choi Date: Mon, 9 Nov 2026 13:27:36 +0900 Subject: [PATCH 5/7] Jet tagging process implemented --- PWGJE/DataModel/JetTagging.h | 144 +++++++++-------- PWGJE/TableProducer/jetGnnInference.cxx | 204 ++++++++++-------------- 2 files changed, 163 insertions(+), 185 deletions(-) diff --git a/PWGJE/DataModel/JetTagging.h b/PWGJE/DataModel/JetTagging.h index e687328f82b..04a12b8fcee 100644 --- a/PWGJE/DataModel/JetTagging.h +++ b/PWGJE/DataModel/JetTagging.h @@ -31,6 +31,7 @@ #include #include #include +#include namespace o2::aod { @@ -169,99 +170,116 @@ 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 of the original track this in-jet track index corresponds to -DECLARE_SOA_COLUMN(TrackIndex, trackIndex, uint8_t); //! in-jet index i (0..N-1) of this constituent track, matching the order of jet.template tracks_as() -DECLARE_SOA_COLUMN(EdgeIndex, edgeIndex, uint8_t); //! unique in-jet edge index for the (i,j) constituent-track pair, i std::pair { - int i = 0; - int offset = 0; - while (offset + (n - 1 - i) <= static_cast(edgeIndex)) { - offset += (n - 1 - i); - i++; - } - int j = i + 1 + (static_cast(edgeIndex) - offset); - return std::pair{static_cast(i), static_cast(j)}; - }); +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_COLUMN(EdgeIndex, edgeIndex, uint8_t); //! unique in-jet edge index for the (i,j) constituent-track pair, i std::pair { +// int i = 0; +// int offset = 0; +// while (offset + (n - 1 - i) <= static_cast(edgeIndex)) { +// offset += (n - 1 - i); +// i++; +// } +// int j = i + 1 + (static_cast(edgeIndex) - offset); +// return std::pair{static_cast(i), static_cast(j)}; +// }); } // namespace jetgnnindex 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(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()) +// 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(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); //! one score per origin class, in jetgnnlabel::Origin order (0..4); producer defines logits/probabilities } // namespace jetgnnscore -#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(_jet_type_##GNNTrack, tracks); \ - } \ - DECLARE_SOA_TABLE(_jet_type_##GNNTrackIndices, "AOD", _description_ "I", \ +// GNNTrack 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 +// 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); -#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); \ - 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); +// #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); \ +// 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); // 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"); DECLARE_JET_GNNTRACKINDEX_TABLE(ChargedMCDetectorLevelJet, "CDJGNT"); -DECLARE_JET_GNNEDGEINDEX_TABLE(ChargedJet, "CJGNE"); -DECLARE_JET_GNNEDGEINDEX_TABLE(ChargedMCDetectorLevelJet, "CDJGNE"); +// DECLARE_JET_GNNEDGEINDEX_TABLE(ChargedJet, "CJGNE"); +// DECLARE_JET_GNNEDGEINDEX_TABLE(ChargedMCDetectorLevelJet, "CDJGNE"); #undef DECLARE_JET_GNNTRACKINDEX_TABLE -#undef DECLARE_JET_GNNEDGEINDEX_TABLE +// #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). DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetTrackOriginDef, "AOD", "CDJTOD", jetgnnlabel::Origin); DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetVertexDef, "AOD", "CDJTVD", jetgnnlabel::VertexIndex); -DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetEdgeDef, "AOD", "CDJEDF", jetgnnlabel::SameVertex); +// DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetEdgeDef, "AOD", "CDJEDF", jetgnnlabel::SameVertex); -// GNN ScoreML outputs, track (node) level, data and MC detector level. Thin tables: join positionally with -// ChargedJetGNNTrack / ChargedMCDetectorLevelJetGNNTrack. -DECLARE_SOA_TABLE(ChargedJetTrackTags, "AOD", "CJTRKT", jetgnnscore::ScoreML); -DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetTrackTags, "AOD", "CDJTRKT", jetgnnscore::ScoreML); +// 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); // 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); - -// Convenience joined types combining the backbone (in-jet ordering / original-track reference) with the thin -// content tables, so `.template tracks_as()` etc. can be used directly. -using ChargedJetGNNTrackLabels = o2::soa::Join; -using ChargedJetGNNEdgeLabels = o2::soa::Join; +// 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()) { +// auto track = prediction.track_as(); +// auto scores = prediction.originScores(); +// } +using ChargedJetGNNTrackPredictions = o2::soa::Join; +using ChargedMCDetectorLevelJetGNNTrackPredictions = o2::soa::Join; +using ChargedJetsWithGNNTrackOrigins = o2::soa::Join; +using ChargedMCDetectorLevelJetsWithGNNTrackOrigins = o2::soa::Join; + +// Convenience joins including truth labels when available (MC detector level). +using ChargedJetGNNTrackLabels = ChargedJetGNNTrackPredictions; +// using ChargedJetGNNEdgeLabels = o2::soa::Join; using ChargedMCDetectorLevelJetGNNTrackLabels = o2::soa::Join; -using ChargedMCDetectorLevelJetGNNEdgeLabels = o2::soa::Join; +// using ChargedMCDetectorLevelJetGNNEdgeLabels = o2::soa::Join; } // namespace o2::aod diff --git a/PWGJE/TableProducer/jetGnnInference.cxx b/PWGJE/TableProducer/jetGnnInference.cxx index e63159336b1..10968cfa652 100644 --- a/PWGJE/TableProducer/jetGnnInference.cxx +++ b/PWGJE/TableProducer/jetGnnInference.cxx @@ -20,12 +20,10 @@ #include "PWGJE/DataModel/JetReducedData.h" #include "PWGJE/DataModel/JetTagging.h" -#include "Common/DataModel/TrackSelectionTables.h" #include "Tools/ML/MlResponse.h" #include #include -#include #include #include #include @@ -37,8 +35,6 @@ #include #include -#include -#include #include #include @@ -46,8 +42,6 @@ #include #include #include -#include -#include #include #include #include @@ -58,36 +52,35 @@ using namespace o2::framework::expressions; namespace jet_gnn_inference { -// Helper struct to define optional produces for JetGnnInferenceTask +// Optional outputs for future track-origin and vertex prediction tables. template struct OptionalProduces : Produces { }; template <> struct OptionalProduces { }; -}; // namespace jet_gnn_inference +} // namespace jet_gnn_inference template struct JetGnnInferenceTask { - static constexpr double DefaultCutsMl[1][2] = {{0.5, 0.5}}; + static constexpr double DefaultCutsMl[1][3] = {{0.5, 0.5, 0.5}}; Produces taggingTable; jet_gnn_inference::OptionalProduces gnnTrackTable; jet_gnn_inference::OptionalProduces gnnVertexTable; - // configuration topological cut for track and sv - Configurable trackDcaXYMax{"trackDcaXYMax", 1, "minimum DCA xy acceptance for tracks [cm]"}; - Configurable trackDcaZMax{"trackDcaZMax", 2, "minimum DCA z acceptance for tracks [cm]"}; + // Track selection for GNN inputs + Configurable trackDcaXYMax{"trackDcaXYMax", 1, "maximum DCA xy acceptance for tracks [cm]"}; + Configurable trackDcaZMax{"trackDcaZMax", 2, "maximum DCA z acceptance for tracks [cm]"}; // ML configuration - Configurable nJetConst{"nJetConst", 100, "maximum number of jet consistuents to be used for ML evaluation"}; + Configurable nJetConst{"nJetConst", 100, "maximum number of jet constituents to be used for ML evaluation"}; Configurable trackPtMin{"trackPtMin", 0.15, "minimum track pT"}; Configurable> binsPtMl{"binsPtMl", std::vector{5., 1000.}, "pT bin limits for ML application"}; - Configurable> cutDirMl{"cutDirMl", std::vector{cuts_ml::CutSmaller, cuts_ml::CutNot}, "Whether to reject score values greater or smaller than the threshold"}; - Configurable> cutsMl{"cutsMl", {DefaultCutsMl[0], 1, 2, {"pT bin 0"}, {"score for default b-jet tagging", "uncer 1"}}, "ML selections per pT bin"}; + Configurable> cutDirMl{"cutDirMl", std::vector{cuts_ml::CutSmaller, cuts_ml::CutNot, cuts_ml::CutNot}, "Whether to reject score values greater or smaller than the threshold"}; + Configurable> cutsMl{"cutsMl", {DefaultCutsMl[0], 1, 3, {"pT bin 0"}, {"b", "c", "lf"}}, "ML selections per pT bin"}; Configurable nClassesMl{"nClassesMl", 3, "Number of classes in ML model"}; - Configurable> namesInputFeatures{"namesInputFeatures", std::vector{"feature1", "feature2"}, "Names of ML model input features"}; Configurable ccdbUrl{"ccdbUrl", "http://alice-ccdb.cern.ch", "url of the ccdb repository"}; Configurable> modelPathsCCDB{"modelPathsCCDB", std::vector{"Users/c/chchoi"}, "Paths of models on CCDB"}; @@ -101,7 +94,7 @@ struct JetGnnInferenceTask { Configurable dbMax{"dbMax", 15., "maximum GNN Db"}; Configurable fC{"fC", 0.018, "Parameter f_c for D_b calculation"}; Configurable nJetFeat{"nJetFeat", 4, "Number of jet GNN input features"}; - Configurable nTrkFeat{"nTrkFeat", 13, "Number of track GNN input features"}; + Configurable nTrkFeat{"nTrkFeat", 8, "Number of track GNN input features"}; Configurable nTrkOrigin{"nTrkOrigin", 5, "Number of track origin categories"}; Configurable> transformFeatureJetMean{"transformFeatureJetMean", std::vector{-999}, @@ -118,59 +111,47 @@ struct JetGnnInferenceTask { Configurable tfFuncTypeGNN{"tfFuncTypeGNN", "asinh", "Transformation function type for GNN"}; o2::analysis::MlResponseHfTagging bMlResponse; - o2::ccdb::CcdbApi ccdbApi; using JetTracksExt = soa::Join; - using OriginalTracks = soa::Join; - - std::vector scoreML; + using AnalysisJets = std::conditional_t, JetTable>; o2::analysis::GNNBjetAllocator tensorAlloc; - template - void fillTables(T const& jet, U const& tracks) + template + void fillHistograms(Jet const& jet, float db) { - int origin = -1; - float jetProb = -1.0; - if constexpr (isMC) { - origin = jet.origin(); + if (jet.pt() < jetpTMin) { + return; } - if (doprocessAlgorithmGNN) { - if (jet.pt() >= jetpTMin) { - float dbRange; - if (scoreML[jet.globalIndex()] < dbMin) { - dbRange = 0.5; // underflow - } else if (scoreML[jet.globalIndex()] < dbMax) { - dbRange = 1.5; // in range - } else { - dbRange = 2.5; // overflow - } - registry.fill(HIST("h2_count_db"), 3.5, dbRange); // incl jet - if constexpr (isMC) { - switch (origin) { - case 2: - registry.fill(HIST("h_db_b"), scoreML[jet.globalIndex()]); - registry.fill(HIST("h2_count_db"), 0.5, dbRange); // b-jet - break; - case 1: - registry.fill(HIST("h_db_c"), scoreML[jet.globalIndex()]); - registry.fill(HIST("h2_count_db"), 1.5, dbRange); // c-jet - break; - case 0: - case 3: - registry.fill(HIST("h_db_lf"), scoreML[jet.globalIndex()]); - registry.fill(HIST("h2_count_db"), 2.5, dbRange); // lf-jet - break; - default: - LOGF(debug, "doprocessAlgorithmGNN, Unexpected origin value: %d (%d)", origin, jet.globalIndex()); - } - } - registry.fill(HIST("h2_pt_db"), jet.pt(), scoreML[jet.globalIndex()]); + float dbRange; + if (db < dbMin) { + dbRange = 0.5; // underflow + } else if (db < dbMax) { + dbRange = 1.5; // in range + } else { + dbRange = 2.5; // overflow + } + registry.fill(HIST("h2_count_db"), 3.5, dbRange); // incl jet + if constexpr (isMCD) { + switch (jet.origin()) { + case 2: + registry.fill(HIST("h_db_b"), db); + registry.fill(HIST("h2_count_db"), 0.5, dbRange); + break; + case 1: + registry.fill(HIST("h_db_c"), db); + registry.fill(HIST("h2_count_db"), 1.5, dbRange); + break; + case 0: + case 3: + registry.fill(HIST("h_db_lf"), db); + registry.fill(HIST("h2_count_db"), 2.5, dbRange); + break; + default: + LOGF(debug, "GNN inference, unexpected origin value: %d (%d)", jet.origin(), jet.globalIndex()); } - // gnnTrackTable(scoreMLAux1[]) - // gnnVertexTable(scoreMLAux2[]) } - taggingTable(0, jetProb, scoreML[jet.globalIndex()]); + registry.fill(HIST("h2_pt_db"), jet.pt(), db); } HistogramRegistry registry{"registry", {}, OutputObjHandlingPolicy::AnalysisObject}; @@ -178,14 +159,24 @@ struct JetGnnInferenceTask { void init(InitContext const&) { if (doprocessAlgorithmGNN) { + // analyzeJetTrackInfo4GNN supplies four jet features and eight track features. + if (nJetFeat != 4 || nTrkFeat != 8 || nClassesMl != 3 || nJetConst <= 0) { + LOG(fatal) << "GNN jet tagging requires nJetFeat=4, nTrkFeat=8, nClassesMl=3 and nJetConst>0."; + } + if (transformFeatureJetMean.value.size() != static_cast(nJetFeat.value) || + transformFeatureJetStdev.value.size() != static_cast(nJetFeat.value) || + transformFeatureTrkMean.value.size() != static_cast(nTrkFeat.value) || + transformFeatureTrkStdev.value.size() != static_cast(nTrkFeat.value)) { + LOG(fatal) << "GNN normalization vectors must match nJetFeat and nTrkFeat."; + } bMlResponse.configure(binsPtMl, cutsMl, cutDirMl, nClassesMl); if (loadModelsFromCCDB) { + o2::ccdb::CcdbApi ccdbApi; ccdbApi.init(ccdbUrl); bMlResponse.setModelPathsCCDB(onnxFileNames, ccdbApi, modelPathsCCDB, timestampCCDB); } else { bMlResponse.setModelPathsLocal(onnxFileNames); } - bMlResponse.cacheInputFeaturesIndices(namesInputFeatures); bMlResponse.init(); tensorAlloc = o2::analysis::GNNBjetAllocator(nJetFeat.value, nTrkFeat.value, nClassesMl.value, nTrkOrigin.value, transformFeatureJetMean.value, transformFeatureJetStdev.value, transformFeatureTrkMean.value, transformFeatureTrkStdev.value, nJetConst, tfFuncTypeGNN.value); @@ -207,39 +198,28 @@ struct JetGnnInferenceTask { } } - template - void analyzeJetAlgorithmGNN(AnyJets const& jets, AnyTracks const& tracks) + template + float evaluateJet(Jet const& jet, Tracks const& tracks) { - for (const auto& jet : jets) { - std::vector> trkFeat; - jettaggingutilities::analyzeJetTrackInfo4GNN(jet, tracks, trkFeat, trackPtMin, trackDcaXYMax, trackDcaZMax, nJetConst); - - std::vector jetFeat{jet.pt(), jet.phi(), jet.eta(), jet.mass()}; - - if (trkFeat.size() > 0) { - std::vector feat; - std::vector gnnInput; - tensorAlloc.getGNNInput(jetFeat, trkFeat, feat, gnnInput); - - auto modelOutput = bMlResponse.getModelOutput(gnnInput, 0); - 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()); - } + std::vector> trkFeat; + jettaggingutilities::analyzeJetTrackInfo4GNN(jet, tracks, trkFeat, trackPtMin, trackDcaXYMax, trackDcaZMax, nJetConst); + if (trkFeat.empty()) { + LOGF(debug, "GNN inference, no selected tracks (%d)", jet.globalIndex()); + return -999.f; } - } - template - void analyzeJetAlgorithmGNNBatch(AnyJets const& jets, AnyTracks const& tracks) - { - // Implement batch inference logic here + std::vector jetFeat{jet.pt(), jet.phi(), jet.eta(), jet.mass()}; + std::vector feat; + std::vector gnnInput; + tensorAlloc.getGNNInput(jetFeat, trkFeat, feat, gnnInput); + + auto modelOutput = bMlResponse.getModelOutput(gnnInput, 0); + float db = jettaggingutilities::getDb(modelOutput, fC); + if (std::isnan(db)) { + LOGF(debug, "GNN inference, Db is NaN (%d)", jet.globalIndex()); + return 999.f; + } + return db; } void processDummy(aod::JetCollisions const&) @@ -247,37 +227,21 @@ struct JetGnnInferenceTask { } PROCESS_SWITCH(JetGnnInferenceTask, processDummy, "Dummy process", true); - void processSetup(JetTable const& jets) + void processAlgorithmGNN(AnalysisJets const& jets, JetTracksExt const& tracks) { - scoreML.clear(); - scoreML.resize(jets.size()); - } - PROCESS_SWITCH(JetGnnInferenceTask, processSetup, "Setup initialization and size of jets for filling table", false); - - void processAlgorithmGNN(JetTable const& jets, JetTracksExt const& jtracks) - { - analyzeJetAlgorithmGNN(jets, jtracks); - } - PROCESS_SWITCH(JetGnnInferenceTask, processAlgorithmGNN, "Fill GNN evaluation score (D_b) for charged jets", false); - - void processAlgorithmGNNBatch(aod::JetCollisions::iterator const& /*coll*/, JetTable const& jets, JetTracksExt const& jtracks) - { - // Batch inference of GNN model per collision - analyzeJetAlgorithmGNNBatch(jets, jtracks); - } - PROCESS_SWITCH(JetGnnInferenceTask, processAlgorithmGNNBatch, "Fill GNN evaluation score (D_b) for charged jets, batch inference", false); - - void processFillTables(std::conditional_t, JetTable>::iterator const& jet, JetTracksExt const& tracks) - { - fillTables(jet, tracks); + // Preserve a one-to-one positional join with Jets, including jets without + // selected tracks and jets below the QA pT threshold. + for (auto const& jet : jets) { + float db = evaluateJet(jet, tracks); + taggingTable(0, -1.f, db); + fillHistograms(jet, db); + } } - PROCESS_SWITCH(JetGnnInferenceTask, processFillTables, "Fill Tables for tagging decision, jet probability, and ML score on charged jets", false); + PROCESS_SWITCH(JetGnnInferenceTask, processAlgorithmGNN, "Run GNN inference and fill jet tagging tables", false); }; using JetGnnTagInferenceDataCharged = JetGnnInferenceTask, aod::ChargedJetTags>; using JetGnnTagInferenceMcdCharged = JetGnnInferenceTask, aod::ChargedMCDetectorLevelJetTags>; -using JetGnnFullInferenceDataCharged = JetGnnInferenceTask, aod::ChargedJetTags, ChargedJetGNNTrack, ChargedJetTrackTags, ChargedJetGNNEdge, ChargedJetEdgeTags>; -using JetGnnFullInferenceMcdCharged = JetGnnInferenceTask, aod::ChargedMCDetectorLevelJetTags, ChargedMCDetectorLevelJetGNNTrack, ChargedMCDetectorLevelJetTrackTags, ChargedMCDetectorLevelJetGNNEdge, ChargedMCDetectorLevelJetEdgeTags>; WorkflowSpec defineDataProcessing(ConfigContext const& cfgc) { @@ -285,9 +249,5 @@ WorkflowSpec defineDataProcessing(ConfigContext const& cfgc) adaptAnalysisTask( cfgc, SetDefaultProcesses{}, TaskName{"jet-gnn-tag-inference-data-charged"}), // o2-linter: disable=name/o2-task (custom device name for templated task alias) adaptAnalysisTask( - cfgc, SetDefaultProcesses{}, TaskName{"jet-gnn-tag-inference-mcd-charged"}), // o2-linter: disable=name/o2-task (custom device name for templated task alias) - adaptAnalysisTask( - cfgc, SetDefaultProcesses{}, TaskName{"jet-gnn-full-inference-data-charged"}), // o2-linter: disable=name/o2-task (custom device name for templated task alias) - adaptAnalysisTask( - cfgc, SetDefaultProcesses{}, TaskName{"jet-gnn-full-inference-mcd-charged"})}; // o2-linter: disable=name/o2-task (custom device name for templated task alias) + cfgc, SetDefaultProcesses{}, TaskName{"jet-gnn-tag-inference-mcd-charged"})}; // o2-linter: disable=name/o2-task (custom device name for templated task alias) } From ad57ddf31532c7ecbba6257a5d31d84c81634f4f Mon Sep 17 00:00:00 2001 From: Changhwan Choi Date: Tue, 10 Nov 2026 15:49:16 +0900 Subject: [PATCH 6/7] Jet GNN inference, track-origin/vertex definition --- PWGJE/Core/JetTaggingUtilities.h | 44 +++- PWGJE/DataModel/JetTagging.h | 120 +++++----- PWGJE/TableProducer/CMakeLists.txt | 5 + PWGJE/TableProducer/jetGnnDefinition.cxx | 258 ++++++++++++++++++++ PWGJE/TableProducer/jetGnnInference.cxx | 284 ++++++++++++++++++++--- PWGJE/Tasks/bjetTaggingGnn.cxx | 50 ++-- PWGJE/Tasks/bjetTreeCreator.cxx | 129 +++++----- 7 files changed, 712 insertions(+), 178 deletions(-) create mode 100644 PWGJE/TableProducer/jetGnnDefinition.cxx diff --git a/PWGJE/Core/JetTaggingUtilities.h b/PWGJE/Core/JetTaggingUtilities.h index 1d07fd40844..b5ed1714d9e 100644 --- a/PWGJE/Core/JetTaggingUtilities.h +++ b/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 -void analyzeJetTrackInfo4GNN(AnalysisJet const& analysisJet, AnyTracks const& /*allTracks*/, std::vector>& 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>& tracksParams, float trackPtMin = 0.5, float trackDcaXYMax = 10.0, float trackDcaZMax = 10.0, int64_t nMaxConstit = 40, std::vector* constituentIndices = nullptr) { - for (const auto& constituent : analysisJet.template tracks_as()) { + // 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(); + 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,6 +1146,9 @@ void analyzeJetTrackInfo4GNN(AnalysisJet const& analysisJet, AnyTracks const& /* if (static_cast(tracksParams.size()) < nMaxConstit) { tracksParams.emplace_back(std::vector{constituent.pt(), constituent.phi(), constituent.eta(), static_cast(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; @@ -1149,8 +1156,12 @@ void analyzeJetTrackInfo4GNN(AnalysisJet const& analysisJet, AnyTracks const& /* 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{constituent.pt(), constituent.phi(), constituent.eta(), static_cast(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& 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 +T getTrackDb(const std::vector& 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 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_ diff --git a/PWGJE/DataModel/JetTagging.h b/PWGJE/DataModel/JetTagging.h index 04a12b8fcee..97f6547938a 100644 --- a/PWGJE/DataModel/JetTagging.h +++ b/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 std::pair { @@ -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); //! 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); \ -// 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); \ + 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()) { -// auto track = prediction.track_as(); -// auto scores = prediction.originScores(); -// } -using ChargedJetGNNTrackPredictions = o2::soa::Join; -using ChargedMCDetectorLevelJetGNNTrackPredictions = o2::soa::Join; -using ChargedJetsWithGNNTrackOrigins = o2::soa::Join; -using ChargedMCDetectorLevelJetsWithGNNTrackOrigins = o2::soa::Join; - -// Convenience joins including truth labels when available (MC detector level). -using ChargedJetGNNTrackLabels = ChargedJetGNNTrackPredictions; -// using ChargedJetGNNEdgeLabels = o2::soa::Join; -using ChargedMCDetectorLevelJetGNNTrackLabels = o2::soa::Join; -// using ChargedMCDetectorLevelJetGNNEdgeLabels = o2::soa::Join; +// ChargedJetGnnEdges / ChargedMCDetectorLevelJetGnnEdges. +// DECLARE_SOA_TABLE(ChargedJetGnnEdgeTags, "AOD", "CJEDGT", jetgnnscore::ScoreML); +// DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetGnnEdgeTags, "AOD", "CDJEDGT", jetgnnscore::ScoreML); } // namespace o2::aod diff --git a/PWGJE/TableProducer/CMakeLists.txt b/PWGJE/TableProducer/CMakeLists.txt index 8f4c523d23f..a865a7e4b78 100644 --- a/PWGJE/TableProducer/CMakeLists.txt +++ b/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 diff --git a/PWGJE/TableProducer/jetGnnDefinition.cxx b/PWGJE/TableProducer/jetGnnDefinition.cxx new file mode 100644 index 00000000000..ce6f6aabc0c --- /dev/null +++ b/PWGJE/TableProducer/jetGnnDefinition.cxx @@ -0,0 +1,258 @@ +// Copyright 2019-2020 CERN and copyright holders of ALICE O2. +// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders. +// All rights not expressly granted are reserved. +// +// This software is distributed under the terms of the GNU General Public +// License v3 (GPL Version 3), copied verbatim in the file "COPYING". +// +// In applying this license CERN does not waive the privileges and immunities +// granted to it by virtue of its status as an Intergovernmental Organization +// or submit itself to any jurisdiction. + +/// \file jetGnnDefinition.cxx +/// \brief Task to produce GNN track indices and MC prediction label definitions +/// +/// \author Changhwan Choi + +#include "PWGJE/Core/JetDerivedDataUtilities.h" +#include "PWGJE/Core/JetTaggingUtilities.h" +#include "PWGJE/DataModel/Jet.h" +#include "PWGJE/DataModel/JetReducedData.h" +#include "PWGJE/DataModel/JetTagging.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace o2; +using namespace o2::framework; + +namespace jet_gnn_definition +{ +// Optional outputs for MC track-origin and vertex definition tables. +template +struct OptionalProduces : Produces { +}; +template <> +struct OptionalProduces { +}; +} // namespace jet_gnn_definition + +template +struct JetGnnDefinitionTask { + static constexpr float largeNegativeNumber = -98.f; + static constexpr float largePositiveNumber = 9999.f; + + Produces> gnnTrackTable; + Produces> gnnTrackIndicesTable; + jet_gnn_definition::OptionalProduces> gnnDefTrackTable; + jet_gnn_definition::OptionalProduces> gnnDefVertexTable; + + Configurable vertexZCut{"vertexZCut", 10.0f, "Accepted z-vertex range"}; + Configurable eventSelections{"eventSelections", "sel8", "choose event selection"}; + std::vector eventSelectionBits; + + // Track selection for GNN inputs + Configurable trackDcaXYMax{"trackDcaXYMax", 1, "maximum DCA xy acceptance for tracks [cm]"}; + Configurable trackDcaZMax{"trackDcaZMax", 2, "maximum DCA z acceptance for tracks [cm]"}; + + // Match the jetFilter and isAcceptedJet criteria in bjetTaggingGnn.cxx. + Configurable jetPtMin{"jetPtMin", 0.0, "minimum jet pT"}; + Configurable jetPtMax{"jetPtMax", 300.0, "maximum jet pT"}; + Configurable jetEtaMin{"jetEtaMin", -0.9, "minimum jet pseudorapidity"}; + Configurable jetEtaMax{"jetEtaMax", 0.9, "maximum jet pseudorapidity"}; + Configurable leadingConstituentPtMin{"leadingConstituentPtMin", -99.0, "minimum pT selection on jet constituent"}; + Configurable leadingConstituentPtMax{"leadingConstituentPtMax", 9999.0, "maximum pT selection on jet constituent"}; + Configurable jetAreaFractionMin{"jetAreaFractionMin", -99.0, "used to make a cut on the jet areas"}; + + Configurable trackPtMin{"trackPtMin", 0.15, "minimum track pT"}; + Configurable nJetConst{"nJetConst", 100, "maximum number of jet constituents to be used for ML evaluation"}; + Configurable vtxRes{"vtxRes", 0.01, "Vertex position resolution (cluster size) for GNN vertex predictions (cm)"}; + Configurable searchUpToQuark{"searchUpToQuark", false, "Flag whether to search for track origins up to the quark level"}; + + using JetTracksExt = std::conditional_t, soa::Join>; + using AnalysisJets = JetTable; + using CollisionsData = aod::JetCollisions; + using CollisionsMcd = soa::Join; + using McCollisions = soa::Join; + using JetParticlesWithIds = soa::Join; + + template + bool isAcceptedCollision(Collision const& collision) + { + if (!(std::abs(collision.posZ()) < vertexZCut) || + !jetderiveddatautilities::selectCollision(collision, eventSelectionBits)) { + return false; + } + if constexpr (isMCD) { + if (collision.isOutlier() || + (collision.has_mcCollision() && collision.template mcCollision_as().isOutlier())) { + return false; + } + } + return true; + } + + template + bool isAcceptedJet(Jet const& jet) + { + // Apply the analysis jetFilter as a predicate rather than dropping rows + // with soa::Filtered: GnnTrackIds must remain joinable to every input jet. + if (!(jet.pt() >= jetPtMin && jet.pt() < jetPtMax && + jet.eta() < jetEtaMax - jet.r() / 100.f && + jet.eta() > jetEtaMin + jet.r() / 100.f)) { + return false; + } + + if (jetAreaFractionMin > largeNegativeNumber && + jet.area() < jetAreaFractionMin * o2::constants::math::PI * (jet.r() / 100.0) * (jet.r() / 100.0)) { + return false; + } + + bool checkConstituentMinPt = leadingConstituentPtMin > largeNegativeNumber; + bool checkConstituentMaxPt = leadingConstituentPtMax < largePositiveNumber; + if (!checkConstituentMinPt && !checkConstituentMaxPt) { + return true; + } + + bool isMinLeadingConstituent = !checkConstituentMinPt; + bool isMaxLeadingConstituent = true; + for (auto const& constituent : jet.template tracks_as()) { + double pt = constituent.pt(); + if (checkConstituentMinPt && pt >= leadingConstituentPtMin) { + isMinLeadingConstituent = true; + } + if (checkConstituentMaxPt && pt > leadingConstituentPtMax) { + isMaxLeadingConstituent = false; + } + } + return isMinLeadingConstituent && isMaxLeadingConstituent; + } + + void init(InitContext const&) + { + eventSelectionBits = jetderiveddatautilities::initialiseEventSelectionBits(eventSelections); + constexpr int MaxGnnTracks = std::numeric_limits::max() + 1; + if ((doprocessMcd || doprocessData) && (nJetConst < 1 || nJetConst > MaxGnnTracks)) { + LOGF(fatal, "nJetConst must be in [1, %d] to fit the uint8_t trackIndex", MaxGnnTracks); + } + if ((doprocessMcd && !isMCD) || (doprocessData && isMCD)) { + LOGF(fatal, "Use processMcd for MC detector-level jets and processData for data jets"); + } + } + + void processDummy(aod::JetCollisions::iterator const&) + { + } + PROCESS_SWITCH(JetGnnDefinitionTask, processDummy, "Dummy process", true); + + template + void fillTrackTables(AnalysisJets const& jets, JetTracksExt const& tracks, bool acceptedCollision, GetTrackLabels const& getTrackLabels) + { + for (auto const& jet : jets) { + // Collision grouping must preserve the positional extension of Jets. + if (gnnTrackIndicesTable.lastIndex() + 1 != jet.globalIndex()) { + LOG(fatal) << "GNN definition requires jets ordered by collision with valid collision indices."; + } + std::vector gnnTrackIds; + if (!acceptedCollision || jet.tracksIds().size() <= 1 || !isAcceptedJet(jet)) { + gnnTrackIndicesTable(gnnTrackIds); + continue; + } + + // Reuse the inference selection, including replacement of low-significance + // tracks at nJetConst. Keep the constituent position for each model node. + std::vector> trackFeatures; + std::vector constituentIndices; + jettaggingutilities::analyzeJetTrackInfo4GNN(jet, tracks, trackFeatures, trackPtMin, trackDcaXYMax, trackDcaZMax, nJetConst, &constituentIndices); + if (constituentIndices.size() <= 1) { + gnnTrackIndicesTable(gnnTrackIds); + continue; + } + + std::vector origins; + std::vector vertexIndices; + if constexpr (isMCD) { + getTrackLabels(jet, origins, vertexIndices); + } + + auto const constituents = jet.template tracks_as(); + gnnTrackIds.reserve(constituentIndices.size()); + for (std::size_t nodeIndex = 0; nodeIndex < constituentIndices.size(); ++nodeIndex) { + auto const constituentIndex = constituentIndices[nodeIndex]; + auto const& track = constituents[constituentIndex]; + gnnTrackTable(jet.globalIndex(), track.globalIndex(), static_cast(nodeIndex)); + if constexpr (isMCD) { + // The backbone and truth table share exactly the same row order. + gnnDefTrackTable(static_cast(origins[constituentIndex])); + if constexpr (!std::is_void_v) { + gnnDefVertexTable(static_cast(vertexIndices[constituentIndex])); + } + } + gnnTrackIds.push_back(gnnTrackTable.lastIndex()); + } + gnnTrackIndicesTable(gnnTrackIds); + } + } + + // PROCESS_SWITCH compiles both entry points for each task specialization, + // including the disabled one. Their table-specific code needs compile-time guards. + void processData(CollisionsData::iterator const& collision, AnalysisJets const& jets, JetTracksExt const& tracks) + { + if constexpr (!isMCD) { + bool acceptedCollision = isAcceptedCollision(collision); + fillTrackTables(jets, tracks, acceptedCollision, [](auto const&, std::vector&, std::vector&) {}); + } + } + PROCESS_SWITCH(JetGnnDefinitionTask, processData, "Produce GNN track indices for data jets", false); + + void processMcd(CollisionsMcd::iterator const& collision, AnalysisJets const& jets, JetTracksExt const& tracks, McCollisions const&, JetParticlesWithIds const& particles, aod::McParticles const& originalParticles) + { + if constexpr (isMCD) { + bool acceptedCollision = isAcceptedCollision(collision); + fillTrackTables(jets, tracks, acceptedCollision, [&](auto const& jet, std::vector& origins, std::vector& vertexIndices) { + // vertexClustering labels the full constituent array, as in the training + // tree creator. Project these labels onto the selected GNN nodes. + std::unordered_map> trackLabels{{"trkVtxIndex", {}}, {"trkOrigin", {}}}; + if (collision.has_mcCollision()) { + jettaggingutilities::vertexClustering(collision.template mcCollision_as(), jet, tracks, particles, originalParticles, trackLabels, searchUpToQuark, vtxRes, trackPtMin); + origins = std::move(trackLabels.at("trkOrigin")); + vertexIndices = std::move(trackLabels.at("trkVtxIndex")); + } else { + // No matched MC collision: unknown origin (0) and no vertex label (-1). + origins.assign(jet.tracksIds().size(), 0); + vertexIndices.assign(jet.tracksIds().size(), -1); + } + }); + } + } + PROCESS_SWITCH(JetGnnDefinitionTask, processMcd, "Produce GNN track indices, origin and vertex definitions for MC detector-level jets", false); +}; + +using JetGnnDefinitionData = JetGnnDefinitionTask>; +using JetGnnDefinitionMcd = JetGnnDefinitionTask, aod::ChargedMCDetectorLevelJetTrackOriginDef, aod::ChargedMCDetectorLevelJetVertexIndexDef>; + +WorkflowSpec defineDataProcessing(ConfigContext const& cfgc) +{ + return WorkflowSpec{ + adaptAnalysisTask( + cfgc, SetDefaultProcesses{}, TaskName{"jet-gnn-definition-data"}), // o2-linter: disable=name/o2-task (custom device name for templated task alias) + adaptAnalysisTask( + cfgc, SetDefaultProcesses{}, TaskName{"jet-gnn-definition-mcd"})}; // o2-linter: disable=name/o2-task (custom device name for templated task alias) +} diff --git a/PWGJE/TableProducer/jetGnnInference.cxx b/PWGJE/TableProducer/jetGnnInference.cxx index 10968cfa652..877404b8624 100644 --- a/PWGJE/TableProducer/jetGnnInference.cxx +++ b/PWGJE/TableProducer/jetGnnInference.cxx @@ -14,6 +14,7 @@ /// /// \author Changhwan Choi +#include "PWGJE/Core/JetDerivedDataUtilities.h" #include "PWGJE/Core/JetTaggingUtilities.h" #include "PWGJE/Core/MlResponseHfTagging.h" #include "PWGJE/DataModel/Jet.h" @@ -23,6 +24,7 @@ #include "Tools/ML/MlResponse.h" #include +#include #include #include #include @@ -48,11 +50,10 @@ using namespace o2; using namespace o2::framework; -using namespace o2::framework::expressions; namespace jet_gnn_inference { -// Optional outputs for future track-origin and vertex prediction tables. +// Optional outputs for track-origin and future vertex prediction tables. template struct OptionalProduces : Produces { }; @@ -61,18 +62,34 @@ struct OptionalProduces { }; } // namespace jet_gnn_inference -template +template struct JetGnnInferenceTask { static constexpr double DefaultCutsMl[1][3] = {{0.5, 0.5, 0.5}}; + static constexpr float invalidGnnScore = -999.f; + static constexpr float largeNegativeNumber = -98.f; + static constexpr float largePositiveNumber = 9999.f; Produces taggingTable; jet_gnn_inference::OptionalProduces gnnTrackTable; jet_gnn_inference::OptionalProduces gnnVertexTable; + Configurable vertexZCut{"vertexZCut", 10.0f, "Accepted z-vertex range"}; + Configurable eventSelections{"eventSelections", "sel8", "choose event selection"}; + std::vector eventSelectionBits; + // Track selection for GNN inputs Configurable trackDcaXYMax{"trackDcaXYMax", 1, "maximum DCA xy acceptance for tracks [cm]"}; Configurable trackDcaZMax{"trackDcaZMax", 2, "maximum DCA z acceptance for tracks [cm]"}; + // Match the jetFilter and isAcceptedJet criteria in bjetTaggingGnn.cxx. + Configurable jetPtMin{"jetPtMin", 0.0, "minimum jet pT"}; + Configurable jetPtMax{"jetPtMax", 300.0, "maximum jet pT"}; + Configurable jetEtaMin{"jetEtaMin", -0.9, "minimum jet pseudorapidity"}; + Configurable jetEtaMax{"jetEtaMax", 0.9, "maximum jet pseudorapidity"}; + Configurable leadingConstituentPtMin{"leadingConstituentPtMin", -99.0, "minimum pT selection on jet constituent"}; + Configurable leadingConstituentPtMax{"leadingConstituentPtMax", 9999.0, "maximum pT selection on jet constituent"}; + Configurable jetAreaFractionMin{"jetAreaFractionMin", -99.0, "used to make a cut on the jet areas"}; + // ML configuration Configurable nJetConst{"nJetConst", 100, "maximum number of jet constituents to be used for ML evaluation"}; Configurable trackPtMin{"trackPtMin", 0.15, "minimum track pT"}; @@ -89,10 +106,11 @@ struct JetGnnInferenceTask { Configurable loadModelsFromCCDB{"loadModelsFromCCDB", false, "Flag to enable or disable the loading of models from CCDB"}; // GNN configuration - Configurable jetpTMin{"jetpTMin", 5., "minimum jet pT"}; + Configurable jetpTMin{"jetpTMin", 5., "minimum jet pT for QA histograms"}; Configurable dbMin{"dbMin", -15., "minimum GNN Db"}; Configurable dbMax{"dbMax", 15., "maximum GNN Db"}; Configurable fC{"fC", 0.018, "Parameter f_c for D_b calculation"}; + Configurable trackFC{"trackFC", 0.018, "Charm fraction for the track-origin D_b calculation"}; Configurable nJetFeat{"nJetFeat", 4, "Number of jet GNN input features"}; Configurable nTrkFeat{"nTrkFeat", 8, "Number of track GNN input features"}; Configurable nTrkOrigin{"nTrkOrigin", 5, "Number of track origin categories"}; @@ -114,11 +132,68 @@ struct JetGnnInferenceTask { using JetTracksExt = soa::Join; using AnalysisJets = std::conditional_t, JetTable>; + using Collisions = std::conditional_t, aod::JetCollisions>; + using McCollisions = soa::Join; + // Data reuses the already subscribed JCollisions table; only MCD requires + // a separate context to dereference the matched MC collision's outlier flag. + using CollisionContext = std::conditional_t; o2::analysis::GNNBjetAllocator tensorAlloc; + template + bool isAcceptedCollision(Collision const& collision) + { + if (!(std::abs(collision.posZ()) < vertexZCut) || + !jetderiveddatautilities::selectCollision(collision, eventSelectionBits)) { + return false; + } + if constexpr (isMCD) { + if (collision.isOutlier() || + (collision.has_mcCollision() && collision.template mcCollision_as().isOutlier())) { + return false; + } + } + return true; + } + + template + bool isAcceptedJet(Jet const& jet) + { + // Apply the analysis jetFilter as a predicate rather than dropping rows + // with soa::Filtered: taggingTable must remain joinable to every input jet. + if (!(jet.pt() >= jetPtMin && jet.pt() < jetPtMax && + jet.eta() < jetEtaMax - jet.r() / 100.f && + jet.eta() > jetEtaMin + jet.r() / 100.f)) { + return false; + } + + if (jetAreaFractionMin > largeNegativeNumber && + jet.area() < jetAreaFractionMin * o2::constants::math::PI * (jet.r() / 100.0) * (jet.r() / 100.0)) { + return false; + } + + bool checkConstituentMinPt = leadingConstituentPtMin > largeNegativeNumber; + bool checkConstituentMaxPt = leadingConstituentPtMax < largePositiveNumber; + if (!checkConstituentMinPt && !checkConstituentMaxPt) { + return true; + } + + bool isMinLeadingConstituent = !checkConstituentMinPt; + bool isMaxLeadingConstituent = true; + for (auto const& constituent : jet.template tracks_as()) { + double pt = constituent.pt(); + if (checkConstituentMinPt && pt >= leadingConstituentPtMin) { + isMinLeadingConstituent = true; + } + if (checkConstituentMaxPt && pt > leadingConstituentPtMax) { + isMaxLeadingConstituent = false; + } + } + return isMinLeadingConstituent && isMaxLeadingConstituent; + } + template - void fillHistograms(Jet const& jet, float db) + void fillHistograms(Jet const& jet, float db, std::vector const* trackScores = nullptr) { if (jet.pt() < jetpTMin) { return; @@ -152,17 +227,68 @@ struct JetGnnInferenceTask { } } registry.fill(HIST("h2_pt_db"), jet.pt(), db); + + // The tag-only jet type has no GNN-track array. Compile this block only + // for full inference, and require MC truth only for the flavour split. + if constexpr (requires { jet.tracksGnnIds(); }) { + if (trackScores) { + auto const gnnTracks = jet.template tracksGnn_as(); + if (trackScores->size() != gnnTracks.size()) { + LOG(fatal) << "Track QA requires one score per GNN backbone row."; + } + for (auto const& track : gnnTracks) { + std::size_t nodeIndex = track.trackIndex(); + if (nodeIndex >= trackScores->size()) { + LOG(fatal) << "GNN trackIndex is out of range for the track QA scores."; + } + float score = (*trackScores)[nodeIndex]; + if (score == invalidGnnScore || std::isnan(score)) { + continue; + } + registry.fill(HIST("h_trackDb"), score); + if constexpr (isMCD) { + switch (track.origin()) { + case 2: + registry.fill(HIST("h_trackDb_b"), score); + break; + case 1: + registry.fill(HIST("h_trackDb_c"), score); + break; + case 0: + case 3: + case 4: + registry.fill(HIST("h_trackDb_lf"), score); + break; + default: + LOGF(debug, "GNN inference, unexpected track origin value: %d", track.origin()); + } + } + } + } + } } HistogramRegistry registry{"registry", {}, OutputObjHandlingPolicy::AnalysisObject}; void init(InitContext const&) { - if (doprocessAlgorithmGNN) { + eventSelectionBits = jetderiveddatautilities::initialiseEventSelectionBits(eventSelections); + if (doprocessGnnTag && doprocessGnnFull) { + LOG(fatal) << "Enable only one of processGnnTag and processGnnFull."; + } + if constexpr (std::is_void_v) { + if (doprocessGnnFull) { + LOG(fatal) << "processGnnFull requires a task with a track prediction output."; + } + } + if (doprocessGnnTag || doprocessGnnFull) { // analyzeJetTrackInfo4GNN supplies four jet features and eight track features. if (nJetFeat != 4 || nTrkFeat != 8 || nClassesMl != 3 || nJetConst <= 0) { LOG(fatal) << "GNN jet tagging requires nJetFeat=4, nTrkFeat=8, nClassesMl=3 and nJetConst>0."; } + if (doprocessGnnFull && (nTrkOrigin != 5 || !std::isfinite(trackFC.value) || trackFC < 0. || trackFC > 1.)) { + LOG(fatal) << "Track-origin inference requires nTrkOrigin=5 and trackFC in [0, 1]."; + } if (transformFeatureJetMean.value.size() != static_cast(nJetFeat.value) || transformFeatureJetStdev.value.size() != static_cast(nJetFeat.value) || transformFeatureTrkMean.value.size() != static_cast(nTrkFeat.value) || @@ -196,16 +322,55 @@ struct JetGnnInferenceTask { registry.add("h_db_lf", "#it{D}_{b} lf-jet;#it{D}_{b}", {HistType::kTH1F, {{50, -20., 30.}}}); registry.add("h2_pt_db", "#it{p}_{T} vs. #it{D}_{b};#it{p}_{T}^{ch jet} (GeV/#it{c}^{2});#it{D}_{b}", {HistType::kTH2F, {{100, 0., 200.}, {50, -20., 30.}}}); } + if (doprocessGnnFull) { + registry.add("h_trackDb", "#it{d}_{b} inclusive track;#it{d}_{b}", {HistType::kTH1F, {{50, -20., 30.}}}); + if constexpr (isMCD) { + registry.add("h_trackDb_b", "#it{d}_{b} b-track;#it{d}_{b}", {HistType::kTH1F, {{50, -20., 30.}}}); + registry.add("h_trackDb_c", "#it{d}_{b} c-track;#it{d}_{b}", {HistType::kTH1F, {{50, -20., 30.}}}); + registry.add("h_trackDb_lf", "#it{d}_{b} lf-track;#it{d}_{b}", {HistType::kTH1F, {{50, -20., 30.}}}); + } + } } template - float evaluateJet(Jet const& jet, Tracks const& tracks) + float evaluateJet(Jet const& jet, Tracks const& tracks, std::vector* trackScores = nullptr) { std::vector> trkFeat; - jettaggingutilities::analyzeJetTrackInfo4GNN(jet, tracks, trkFeat, trackPtMin, trackDcaXYMax, trackDcaZMax, nJetConst); - if (trkFeat.empty()) { - LOGF(debug, "GNN inference, no selected tracks (%d)", jet.globalIndex()); - return -999.f; + if (trackScores) { + trackScores->clear(); + if constexpr (!std::is_void_v && requires { jet.tracksGnnIds(); }) { + // The definition backbone owns track selection and model-node order. + // Do not reapply inference-side cuts or select a different track subset. + auto const gnnTracks = jet.template tracksGnn_as(); + trackScores->assign(gnnTracks.size(), invalidGnnScore); + if (gnnTracks.size() <= 1) { + return invalidGnnScore; + } + if (gnnTracks.size() > static_cast(nJetConst.value)) { + LOG(fatal) << "GNN backbone contains " << gnnTracks.size() << " tracks, exceeding nJetConst=" << nJetConst.value << ". Increase the inference allocator limit to match jetGnnDefinition."; + } + trkFeat.reserve(gnnTracks.size()); + for (std::size_t nodeIndex = 0; nodeIndex < gnnTracks.size(); ++nodeIndex) { + auto const& node = gnnTracks[nodeIndex]; + if (node.trackIndex() != nodeIndex || node.jetId() != jet.globalIndex()) { + LOG(fatal) << "GNN backbone node order or jet reference does not match the jet's tracksGnnIds()."; + } + auto const track = node.template track_as(); + int sign = jettaggingutilities::getGeoSign(jet, track); + trkFeat.emplace_back(std::vector{track.pt(), track.phi(), track.eta(), static_cast(track.sign()), std::abs(track.dcaXY()) * sign, track.sigmadcaXY(), std::abs(track.dcaZ()) * sign, track.sigmadcaZ()}); + } + } else { + LOG(fatal) << "Track-origin inference requires a task with a track prediction output."; + } + } else { + if (jet.tracksIds().size() <= 1 || !isAcceptedJet(jet)) { + return invalidGnnScore; + } + jettaggingutilities::analyzeJetTrackInfo4GNN(jet, tracks, trkFeat, trackPtMin, trackDcaXYMax, trackDcaZMax, nJetConst); + if (trkFeat.size() <= 1) { + LOGF(debug, "GNN inference, fewer than two selected tracks (%d)", jet.globalIndex()); + return invalidGnnScore; + } } std::vector jetFeat{jet.pt(), jet.phi(), jet.eta(), jet.mass()}; @@ -213,8 +378,33 @@ struct JetGnnInferenceTask { std::vector gnnInput; tensorAlloc.getGNNInput(jetFeat, trkFeat, feat, gnnInput); - auto modelOutput = bMlResponse.getModelOutput(gnnInput, 0); - float db = jettaggingutilities::getDb(modelOutput, fC); + float db; + if (trackScores) { + // Exported head order: edge_outputs, node_outputs, flav_preds. + // Retrieve jet and track logits in one model evaluation; leave edges unused. + auto const outputs = bMlResponse.getModelOutputs(gnnInput, 0); + if (outputs.empty()) { + LOG(error) << "GNN returned no outputs; keeping invalid jet and track scores."; + return invalidGnnScore; + } + if (outputs.size() != 3 || outputs.back().size() != static_cast(nClassesMl.value) || + outputs[1].size() != trkFeat.size() * static_cast(nTrkOrigin.value)) { + LOG(fatal) << "Expected GNN outputs [edges, N*5 track-origin logits, 3 jet logits]."; + } + db = jettaggingutilities::getDb(outputs.back(), fC); + auto const& nodeLogits = outputs[1]; + std::vector trackLogits(nTrkOrigin.value); + for (std::size_t nodeIndex = 0; nodeIndex < trkFeat.size(); ++nodeIndex) { + for (std::size_t category = 0; category < trackLogits.size(); ++category) { + trackLogits[category] = nodeLogits[nodeIndex * trackLogits.size() + category]; + } + float score = jettaggingutilities::getTrackDb(trackLogits, trackFC); + (*trackScores)[nodeIndex] = std::isnan(score) ? invalidGnnScore : score; + } + } else { + auto const modelOutput = bMlResponse.getModelOutput(gnnInput, 0); + db = jettaggingutilities::getDb(modelOutput, fC); + } if (std::isnan(db)) { LOGF(debug, "GNN inference, Db is NaN (%d)", jet.globalIndex()); return 999.f; @@ -222,32 +412,74 @@ struct JetGnnInferenceTask { return db; } - void processDummy(aod::JetCollisions const&) + void processDummy(aod::JetCollisions::iterator const&) { } PROCESS_SWITCH(JetGnnInferenceTask, processDummy, "Dummy process", true); - void processAlgorithmGNN(AnalysisJets const& jets, JetTracksExt const& tracks) + void processGnnTag(typename Collisions::iterator const& collision, AnalysisJets const& jets, JetTracksExt const& tracks, CollisionContext const&) { - // Preserve a one-to-one positional join with Jets, including jets without - // selected tracks and jets below the QA pT threshold. + bool acceptedCollision = isAcceptedCollision(collision); + // Preserve a one-to-one positional join with Jets, including jets rejected + // by collision/jet selection or with fewer than two GNN input tracks. for (auto const& jet : jets) { - float db = evaluateJet(jet, tracks); + if (taggingTable.lastIndex() + 1 != jet.globalIndex()) { + LOG(fatal) << "GNN inference requires jets ordered by collision with valid collision indices."; + } + float db = acceptedCollision ? evaluateJet(jet, tracks) : invalidGnnScore; taggingTable(0, -1.f, db); - fillHistograms(jet, db); + if (acceptedCollision) { + fillHistograms(jet, db); + } + } + } + PROCESS_SWITCH(JetGnnInferenceTask, processGnnTag, "Run GNN inference and fill jet tagging tables", false); + + void processGnnFull(typename Collisions::iterator const& collision, soa::Join const& jets, JetTracksExt const& tracks, JetGnnTracksTable const&, CollisionContext const&) + { + if constexpr (!std::is_void_v) { + bool acceptedCollision = isAcceptedCollision(collision); + std::vector trackScores; + for (auto const& jet : jets) { + if (taggingTable.lastIndex() + 1 != jet.globalIndex()) { + LOG(fatal) << "GNN inference requires jets ordered by collision with valid collision indices."; + } + float db = invalidGnnScore; + if (acceptedCollision) { + db = evaluateJet(jet, tracks, &trackScores); + } else { + // A different definition-side selection may have left backbone rows. + // Keep one invalid prediction per row instead of breaking the join. + trackScores.assign(jet.tracksGnnIds().size(), invalidGnnScore); + } + taggingTable(0, -1.f, db); + auto const nodes = jet.template tracksGnn_as(); + for (std::size_t nodeIndex = 0; nodeIndex < nodes.size(); ++nodeIndex) { + // Thin prediction tables join positionally with the input backbone. + if (nodes[nodeIndex].globalIndex() != gnnTrackTable.lastIndex() + 1) { + LOG(fatal) << "GNN backbone rows are not in jet/node order; refusing a misaligned track prediction table."; + } + gnnTrackTable(trackScores[nodeIndex]); + } + if (acceptedCollision) { + fillHistograms(jet, db, &trackScores); + } + } } } - PROCESS_SWITCH(JetGnnInferenceTask, processAlgorithmGNN, "Run GNN inference and fill jet tagging tables", false); + PROCESS_SWITCH(JetGnnInferenceTask, processGnnFull, "Run GNN inference and fill jet tagging and track-origin prediction tables", false); }; -using JetGnnTagInferenceDataCharged = JetGnnInferenceTask, aod::ChargedJetTags>; -using JetGnnTagInferenceMcdCharged = JetGnnInferenceTask, aod::ChargedMCDetectorLevelJetTags>; +// One task per jet type owns both outputs. Process switches select tag-only or +// full inference without registering duplicate producers of the jet tags. +using JetGnnInferenceData = JetGnnInferenceTask, aod::ChargedJetTags, aod::ChargedJetGnnTrackIds, aod::ChargedJetGnnTracks, aod::ChargedJetGnnTrackOriginTags>; +using JetGnnInferenceMcd = JetGnnInferenceTask, aod::ChargedMCDetectorLevelJetTags, aod::ChargedMCDetectorLevelJetGnnTrackIds, soa::Join, aod::ChargedMCDetectorLevelJetGnnTrackOriginTags>; WorkflowSpec defineDataProcessing(ConfigContext const& cfgc) { return WorkflowSpec{ - adaptAnalysisTask( - cfgc, SetDefaultProcesses{}, TaskName{"jet-gnn-tag-inference-data-charged"}), // o2-linter: disable=name/o2-task (custom device name for templated task alias) - adaptAnalysisTask( - cfgc, SetDefaultProcesses{}, TaskName{"jet-gnn-tag-inference-mcd-charged"})}; // o2-linter: disable=name/o2-task (custom device name for templated task alias) + adaptAnalysisTask( + cfgc, SetDefaultProcesses{}, TaskName{"jet-gnn-inference-data"}), // o2-linter: disable=name/o2-task (custom device name for templated task alias) + adaptAnalysisTask( + cfgc, SetDefaultProcesses{}, TaskName{"jet-gnn-inference-mcd"})}; // o2-linter: disable=name/o2-task (custom device name for templated task alias) } diff --git a/PWGJE/Tasks/bjetTaggingGnn.cxx b/PWGJE/Tasks/bjetTaggingGnn.cxx index 35ace3cc4f4..797f562983e 100644 --- a/PWGJE/Tasks/bjetTaggingGnn.cxx +++ b/PWGJE/Tasks/bjetTaggingGnn.cxx @@ -1191,26 +1191,22 @@ struct BjetTaggingGnn { return false; } - // Shared entry checks for MCD collisions, used by processMCDJets(SV) and processMCDTracks: - // event selection, MC-outlier rejection (own and matched mcCollision's), and eval/train split. - template - bool isRejectedMCDCollision(AnalysisCollision const& collision) + // Physical collision selection matches the GNN definition/inference tasks. + // Keep dataset splitting separate in the MCD process entry checks. + template + bool isAcceptedCollision(Collision const& collision) { - if (!jetderiveddatautilities::selectCollision(collision, eventSelectionBits)) { - return true; + if (!(std::abs(collision.posZ()) < vertexZCut) || + !jetderiveddatautilities::selectCollision(collision, eventSelectionBits)) { + return false; } - // Reject outlier MC collisions - if (collision.isOutlier()) { - return true; - } - if (collision.has_mcCollision() && collision.template mcCollision_as().isOutlier()) { - return true; - } - // Uses only collisionId % trainingDatasetRaioParam != 0 for evaluation dataset - if (trainingDatasetRatioParam && collision.collisionId() % trainingDatasetRatioParam == 0) { - return true; + if constexpr (isMCD) { + if (collision.isOutlier() || + (collision.has_mcCollision() && collision.template mcCollision_as().isOutlier())) { + return false; + } } - return false; + return true; } template @@ -1345,7 +1341,7 @@ struct BjetTaggingGnn { template void processDataJetsCommon(FilteredCollisions::iterator const& collision, AnalysisJets const& alljets, FilteredTracks const& allTracks, SecondaryVertices const* allSVs = nullptr) { - if (!jetderiveddatautilities::selectCollision(collision, eventSelectionBits)) { + if (!isAcceptedCollision(collision)) { return; } // Keeps hCollCounter populated consistently across every process[Data,MCD]Jets(SV)(Sub) variant, not @@ -1405,7 +1401,7 @@ struct BjetTaggingGnn { return; } - if (!jetderiveddatautilities::selectCollision(collision, eventSelectionBits)) { + if (!isAcceptedCollision(collision)) { return; } registry.fill(HIST("h_event_counter"), 0.0); // Coll+TVX+NoTFB+NoITSROFB+Zvtx @@ -1435,7 +1431,7 @@ struct BjetTaggingGnn { return; } - if (!jetderiveddatautilities::selectCollision(collision, eventSelectionBits)) { + if (!isAcceptedCollision(collision)) { return; } registry.fill(HIST("h_event_counter"), 0.0); // Coll+TVX+NoTFB+NoITSROFB+Zvtx @@ -1487,7 +1483,7 @@ struct BjetTaggingGnn { void processDataTracks(FilteredCollisions::iterator const& collision, AnalysisTracks const& tracks) { - if (!jetderiveddatautilities::selectCollision(collision, eventSelectionBits)) { + if (!isAcceptedCollision(collision)) { return; } @@ -1518,7 +1514,11 @@ struct BjetTaggingGnn { template void processMCDJetsCommon(FilteredCollisionsMCD::iterator const& collision, AnalysisJets const& MCDjets, FilteredTracksMCD const& allTracks, FilteredMCPJets const& MCPjets, SecondaryVertices const* allSVs = nullptr) { - if (isRejectedMCDCollision(collision)) { + if (!isAcceptedCollision(collision)) { + return; + } + // Uses only collisionId % trainingDatasetRatioParam != 0 for evaluation. + if (trainingDatasetRatioParam && collision.collisionId() % trainingDatasetRatioParam == 0) { return; } bool matchedMcColl = collision.has_mcCollision(); @@ -1896,7 +1896,11 @@ struct BjetTaggingGnn { void processMCDTracks(FilteredCollisionsMCD::iterator const& collision, AnalysisTracksMCD const& tracks, AnalysisCollisionsMCP const& /*mcCollisions*/, aod::JetParticles const& allParticles) { - if (isRejectedMCDCollision(collision)) { + if (!isAcceptedCollision(collision)) { + return; + } + // Uses only collisionId % trainingDatasetRatioParam != 0 for evaluation. + if (trainingDatasetRatioParam && collision.collisionId() % trainingDatasetRatioParam == 0) { return; } bool matchedMcColl = collision.has_mcCollision(); diff --git a/PWGJE/Tasks/bjetTreeCreator.cxx b/PWGJE/Tasks/bjetTreeCreator.cxx index 186842c0153..0f9251e05e6 100644 --- a/PWGJE/Tasks/bjetTreeCreator.cxx +++ b/PWGJE/Tasks/bjetTreeCreator.cxx @@ -34,6 +34,7 @@ #include #include #include +#include #include #include @@ -46,7 +47,7 @@ #include #include #include -#include +#include #include #include @@ -278,11 +279,7 @@ struct BJetTreeCreator { Configurable produceSVTree{"produceSVTree", true, "produce the SV-correlated jet TTree”"}; Configurable produceTree{"produceTree", true, "produce the jet TTree"}; - Configurable vtxRes{"vtxRes", 0.01, "Vertex position resolution (cluster size) for GNN vertex predictions (cm)"}; - Configurable searchUpToQuark{"searchUpToQuark", true, "Flag whether to search for track origins up to the quark level"}; - Configurable trainingDatasetRatioParam{"trainingDatasetRatioParam", 0, "Parameter for splitting training/evaluation datasets by collisionId"}; - Configurable produceExtrbTableGnn{"produceExtrbTableGnn", false, "Flag whether to produce the extra table with track quality variables for GNN training"}; std::vector eventSelectionBits; @@ -292,6 +289,10 @@ struct BJetTreeCreator { void init(InitContext const&) { + if (doprocessMCJetsForGNN && doprocessMCJetsForGNNExtrb) { + LOGF(fatal, "Enable only one of processMCJetsForGNN and processMCJetsForGNNExtrb"); + } + // Seed the random number generator using current time std::srand(static_cast(std::time(nullptr))); @@ -313,7 +314,7 @@ struct BJetTreeCreator { registry.add("h2_jetMass_jetpT", "Jet mass;#it{p}_{T,jet} (GeV/#it{c});#it{m}_{jet} (GeV/#it{c}^{2})", {HistType::kTH2F, {{200, 0., 200.}, {50, 0, 50.0}}}); registry.add("h2_SVMass_jetpT", "Secondary vertex mass;#it{p}_{T,jet} (GeV/#it{c});#it{m}_{SV} (GeV/#it{c}^{2})", {HistType::kTH2F, {{200, 0., 200.}, {50, 0, 10}}}); - if (doprocessMCJets || doprocessMCJetsForGNN) { + if (doprocessMCJets || doprocessMCJetsForGNN || doprocessMCJetsForGNNExtrb) { registry.add("h2_SIPs2D_jetpT_bjet", "2D IP significance b-jets;#it{p}_{T,jet} (GeV/#it{c});IPs", {HistType::kTH2F, {{200, 0., 200.}, {100, -50.0, 50.0}}}); registry.add("h2_SIPs3D_jetpT_bjet", "3D IP significance b-jets;#it{p}_{T,jet} (GeV/#it{c});IPs", {HistType::kTH2F, {{200, 0., 200.}, {100, -50.0, 50.0}}}); registry.add("h2_LxyS_jetpT_bjet", "Decay length in XY b-jets;#it{p}_{T,jet} (GeV/#it{c});S#it{L}_{xy}", {HistType::kTH2F, {{200, 0., 200.}, {100, 0., 100.0}}}); @@ -348,7 +349,7 @@ struct BJetTreeCreator { registry.add("h2_Response_DetjetpT_PartjetpT_lfjet", "Response matrix lf-jet;#it{p}_{T,jet}^{det} (GeV/#it{c});#it{p}_{T,jet}^{part} (GeV/#it{c})", {HistType::kTH2F, {{200, 0., 200.}, {200, 0., 200.}}}); } - if (doprocessMCJetsForGNN) { + if (doprocessMCJetsForGNN || doprocessMCJetsForGNNExtrb) { //+jet registry.add("h_jet_pt", "jet_pt;#it{p}_{T}^{ch jet} (GeV/#it{c});Entries", {HistType::kTH1F, {{200, 0., 200.}}}); registry.add("h_jet_eta", "jet_eta;#it{#eta}_{ch jet};Entries", {HistType::kTH1F, {{200, -2., 2.}}}); @@ -371,7 +372,7 @@ struct BJetTreeCreator { registry.add("h_trk_dcaz", "trk_dcaxyz;#it{DCA}_{z} (cm);Entries", {HistType::kTH1F, {{200, -0.1, 0.1}}}); registry.add("h_trk_sigmadcaxy", "trk_sigmadcaxy;#it{#sigma}_{#it{DCA}_{xy}} (cm);Entries", {HistType::kTH1F, {{200, 0., 0.1}}}); registry.add("h_trk_sigmadcaz", "trk_sigmadcaxyz;#it{#sigma}_{#it{DCA}_{z}} (cm);Entries", {HistType::kTH1F, {{200, 0., 0.1}}}); - if (produceExtrbTableGnn) { + if (doprocessMCJetsForGNNExtrb) { registry.add("h_trk_itsncls", "trk_itsncls;ITS NCls;Entries", {HistType::kTH1F, {{10, 0., 10.}}}); registry.add("h_trk_tpcncls", "trk_tpcncls;TPC NCls (Found);Entries", {HistType::kTH1F, {{200, 0., 200.}}}); registry.add("h_trk_tpcncrs", "trk_tpcncrs;TPC NCrossedRows;Entries", {HistType::kTH1F, {{200, 0., 200.}}}); @@ -399,6 +400,7 @@ struct BJetTreeCreator { using FilteredCollision = soa::Filtered>; using JetTrackswID = soa::Filtered>; using JetTracksMCDwID = soa::Filtered>; + using JetTracksGNN = soa::Filtered>; using DataJets = soa::Filtered>; using OriginalTracks = soa::Join; @@ -514,8 +516,6 @@ struct BJetTreeCreator { return true; } - using TrackLabelMap = std::unordered_map>; - template void analyzeJetTrackInfo(AnyCollision const& /*collision*/, AnalysisJet const& analysisJet, AnyTracks const& /*allTracks*/, SecondaryVertices const& /*allSVs*/, std::vector& trackIndices, int jetFlavor = 0, double eventweight = 1.0) { @@ -561,21 +561,25 @@ struct BJetTreeCreator { } } - template - void analyzeJetTrackInfoForGNN(AnyCollision const& /*collision*/, AnalysisJet const& analysisJet, AnyTracks const& /*allTracks*/, AnyOriginalTracks const&, std::vector& trackIndices, int jetFlavor = 0, double eventweight = 1.0, TrackLabelMap* trkLabels = nullptr) + template + int analyzeJetTrackInfoForGNN(AnalysisJet const& analysisJet, AnyTracks const& /*allTracks*/, std::vector& trackIndices, int jetFlavor = 0, double eventweight = 1.0) { - int trkIdx = -1; - for (const auto& constituent : analysisJet.template tracks_as()) { - - trkIdx++; + std::unordered_set vertexIndices; + for (auto const& definition : analysisJet.template tracksGnn_as>()) { + auto const constituent = definition.template track_as(); - // No DCA-acceptance cut here: matches BjetTaggingGnn::fillMCDJetHistograms(), which only applies - // a track pT cut (trackPtMinGnn) for GNN input tracks, so the training tree stays consistent with - // what is evaluated at inference time. + // Keep the tree's pT cut on the GNN tracks selected by jetGnnDefinition. + // The backbone supplies both track identity and model-node ordering. if (constituent.pt() < trackPtMin) { continue; } + int trkVtxIndex = definition.vertexIndex(); + int trkOrigin = definition.origin(); + if (trkVtxIndex >= 0) { + vertexIndices.insert(trkVtxIndex); + } + double deltaRJetTrack = jetutilities::deltaR(analysisJet, constituent); double dotProduct = RecoDecay::dotProd(std::array{analysisJet.px(), analysisJet.py(), analysisJet.pz()}, std::array{constituent.px(), constituent.py(), constituent.pz()}); int sign = jettaggingutilities::getGeoSign(analysisJet, constituent); @@ -583,7 +587,7 @@ struct BJetTreeCreator { registry.fill(HIST("h2_SIPs2D_jetpT"), analysisJet.pt(), sign * std::abs(constituent.dcaXY()) / constituent.sigmadcaXY(), eventweight); registry.fill(HIST("h2_SIPs3D_jetpT"), analysisJet.pt(), sign * std::abs(constituent.dcaXYZ()) / constituent.sigmadcaXYZ(), eventweight); - if (doprocessMCJetsForGNN) { + if (doprocessMCJetsForGNN || doprocessMCJetsForGNNExtrb) { if (jetFlavor == JetTaggingSpecies::beauty) { registry.fill(HIST("h2_SIPs2D_jetpT_bjet"), analysisJet.pt(), sign * std::abs(constituent.dcaXY()) / constituent.sigmadcaXY(), eventweight); registry.fill(HIST("h2_SIPs3D_jetpT_bjet"), analysisJet.pt(), sign * std::abs(constituent.dcaXYZ()) / constituent.sigmadcaXYZ(), eventweight); @@ -595,16 +599,6 @@ struct BJetTreeCreator { registry.fill(HIST("h2_SIPs3D_jetpT_lfjet"), analysisJet.pt(), sign * std::abs(constituent.dcaXYZ()) / constituent.sigmadcaXYZ(), eventweight); } - auto origConstit = constituent.template track_as(); - - //+ - int trkVtxIndex = 0; - int trkOrigin = 0; - if (trkLabels != nullptr) { - trkVtxIndex = (*trkLabels)["trkVtxIndex"][trkIdx]; - trkOrigin = (*trkLabels)["trkOrigin"][trkIdx]; - } - //+trk registry.fill(HIST("h_trk_pt"), constituent.pt(), eventweight); registry.fill(HIST("h_trk_eta"), constituent.eta(), eventweight); @@ -614,21 +608,22 @@ struct BJetTreeCreator { registry.fill(HIST("h_trk_dcaz"), std::abs(constituent.dcaZ()) * sign, eventweight); registry.fill(HIST("h_trk_sigmadcaxy"), constituent.sigmadcaXY(), eventweight); registry.fill(HIST("h_trk_sigmadcaz"), constituent.sigmadcaZ(), eventweight); - if (produceExtrbTableGnn) { + registry.fill(HIST("h_trk_vtx_index"), trkVtxIndex); + registry.fill(HIST("h_trk_origin"), trkOrigin); + + if constexpr (withExtrb) { + auto const origConstit = constituent.template track_as(); registry.fill(HIST("h_trk_itsncls"), origConstit.itsNCls(), eventweight); registry.fill(HIST("h_trk_tpcncls"), origConstit.tpcNClsFound(), eventweight); registry.fill(HIST("h_trk_tpcncrs"), origConstit.tpcNClsCrossedRows(), eventweight); registry.fill(HIST("h_trk_itschi2ncl"), origConstit.itsChi2NCl(), eventweight); registry.fill(HIST("h_trk_tpcchi2ncl"), origConstit.tpcChi2NCl(), eventweight); registry.fill(HIST("h2_trk_jtrackpt_vs_origtrackpt"), constituent.pt(), origConstit.pt(), eventweight); // jtrack & new extra table are well joined (linear correlation) - } - registry.fill(HIST("h_trk_vtx_index"), trkVtxIndex); - registry.fill(HIST("h_trk_origin"), trkOrigin); - - if (produceTree) { - if (produceExtrbTableGnn) { + if (produceTree) { bjetTracksExtrbTable(/*bjetParamsTable.lastIndex() + 1, */ constituent.phi(), constituent.sign(), origConstit.itsChi2NCl(), origConstit.tpcChi2NCl(), origConstit.itsNCls(), origConstit.tpcNClsFound(), origConstit.tpcNClsCrossedRows(), trkOrigin, trkVtxIndex); //+ - } else { + } + } else { + if (produceTree) { bjetTracksExtraTable(/*bjetParamsTable.lastIndex() + 1, */ constituent.phi(), constituent.sign(), trkOrigin, trkVtxIndex); //+ } } @@ -639,6 +634,9 @@ struct BJetTreeCreator { } trackIndices.push_back(bjetTracksParamsTable.lastIndex()); } + // Full-jet labels can have gaps after GNN selection and the tree pT cut. + // Count distinct vertices represented by the stored tracks, including PV. + return static_cast(vertexIndices.size()); } void processDummy(FilteredCollision::iterator const& /*collision*/) @@ -800,27 +798,33 @@ struct BJetTreeCreator { PROCESS_SWITCH(BJetTreeCreator, processMCJets, "jet information in MC", false); using FilteredCollisionMCDOutlier = soa::Filtered>; - using MCDJetTableNoSV = soa::Filtered>; using AnalysisCollisionsMCPOutlier = soa::Join; - using JetParticleswID = soa::Join; - // aod::McParticles::vx(), vy(), vz() are necessary for jettaggingutilities::vertexClustering() - void processMCJetsForGNN(FilteredCollisionMCDOutlier::iterator const& collision, MCDJetTableNoSV const& MCDjets, MCPJetTable const& MCPjets, JetTracksMCDwID const& allTracks, JetParticleswID const& MCParticles, AnalysisCollisionsMCPOutlier const&, OriginalTracks const& origTracks, aod::McParticles const& origParticles) + template + bool isAcceptedCollision(Collision const& collision) { - if (!jetderiveddatautilities::selectCollision(collision, eventSelectionBits) || (static_cast(std::rand()) / RAND_MAX < eventReductionFactor)) { - return; + if (!(std::abs(collision.posZ()) < vertexZCut) || + !jetderiveddatautilities::selectCollision(collision, eventSelectionBits)) { + return false; } - - // Uses only collisionId % trainingDatasetRaioParam == 0 for training dataset - if (trainingDatasetRatioParam != 0 && collision.collisionId() % trainingDatasetRatioParam != 0) { - return; + if constexpr (isMCD) { + if (collision.isOutlier() || + (collision.has_mcCollision() && collision.template mcCollision_as().isOutlier())) { + return false; + } } + return true; + } - // Reject outlier MC collisions - if (collision.isOutlier()) { + template + void fillMCJetsForGNN(FilteredCollisionMCDOutlier::iterator const& collision, soa::Filtered> const& MCDjets, AnyTracks const& allTracks) + { + if (!isAcceptedCollision(collision) || (static_cast(std::rand()) / RAND_MAX < eventReductionFactor)) { return; } - if (collision.has_mcCollision() && collision.template mcCollision_as().isOutlier()) { + + // Uses only collisionId % trainingDatasetRaioParam == 0 for training dataset + if (trainingDatasetRatioParam != 0 && collision.collisionId() % trainingDatasetRatioParam != 0) { return; } @@ -828,11 +832,8 @@ struct BJetTreeCreator { registry.fill(HIST("h_vertexZ"), collision.posZ(), eventWeight); - auto const mcParticlesPerColl = MCParticles.sliceBy(mcParticlesPerCollision, collision.mcCollisionId()); - auto const mcPJetsPerColl = MCPjets.sliceBy(mcpJetsPerCollision, collision.mcCollisionId()); - for (const auto& analysisJet : MCDjets) { - if (!isAcceptedJet(analysisJet)) { + if (!isAcceptedJet(analysisJet) || !analysisJet.has_tracksGnn()) { continue; } @@ -845,10 +846,7 @@ struct BJetTreeCreator { continue; } - //+ - TrackLabelMap trkLabels{{"trkVtxIndex", {}}, {"trkOrigin", {}}}; - int nVertices = jettaggingutilities::vertexClustering(collision.template mcCollision_as(), analysisJet, allTracks, MCParticles, origParticles, trkLabels, searchUpToQuark, vtxRes, trackPtMin); - analyzeJetTrackInfoForGNN(collision, analysisJet, allTracks, origTracks, indicesTracks, jetFlavor, eventWeight, &trkLabels); + int nVertices = analyzeJetTrackInfoForGNN(analysisJet, allTracks, indicesTracks, jetFlavor, eventWeight); registry.fill(HIST("h2_jetMass_jetpT"), analysisJet.pt(), analysisJet.mass(), eventWeight); registry.fill(HIST("h2_nTracks_jetpT"), analysisJet.pt(), indicesTracks.size(), eventWeight); @@ -880,7 +878,18 @@ struct BJetTreeCreator { } } } - PROCESS_SWITCH(BJetTreeCreator, processMCJetsForGNN, "jet information in MC for GNN", false); + + void processMCJetsForGNN(FilteredCollisionMCDOutlier::iterator const& collision, soa::Filtered> const& MCDjets, JetTracksGNN const& allTracks, soa::Join const&, AnalysisCollisionsMCPOutlier const&) + { + fillMCJetsForGNN(collision, MCDjets, allTracks); + } + PROCESS_SWITCH(BJetTreeCreator, processMCJetsForGNN, "jet information in MC for GNN using derived tracks", false); + + void processMCJetsForGNNExtrb(FilteredCollisionMCDOutlier::iterator const& collision, soa::Filtered> const& MCDjets, JetTrackswID const& allTracks, soa::Join const&, AnalysisCollisionsMCPOutlier const&, OriginalTracks const&) + { + fillMCJetsForGNN(collision, MCDjets, allTracks); + } + PROCESS_SWITCH(BJetTreeCreator, processMCJetsForGNNExtrb, "jet information in MC for GNN with original-track quality variables", false); Filter mccollisionFilter = nabs(aod::jmccollision::posZ) < vertexZCut; using FilteredCollisionMCP = soa::Filtered; From ba59beafb562e9d29b6b949e45dfb00ba2939b59 Mon Sep 17 00:00:00 2001 From: Changhwan Choi Date: Tue, 10 Nov 2026 20:48:18 +0900 Subject: [PATCH 7/7] Resolved MegaLinter errors --- Tools/ML/model.h | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/Tools/ML/model.h b/Tools/ML/model.h index 00d3d2d552a..02109b6a8a9 100644 --- a/Tools/ML/model.h +++ b/Tools/ML/model.h @@ -24,6 +24,7 @@ #include +#include #include #include #include @@ -230,9 +231,9 @@ class OnnxModel return mSession; } int getNumInputNodes() const { return mInputShapes[0][1]; } - std::vector> getInputShapes() const { return mInputShapes; } + const std::vector>& getInputShapes() const { return mInputShapes; } int getNumOutputNodes() const { return mOutputShapes[0][1]; } - std::vector> getOutputShapes() const { return mOutputShapes; } + const std::vector>& getOutputShapes() const { return mOutputShapes; } uint64_t getValidityFrom() const { return validFrom; } uint64_t getValidityUntil() const { return validUntil; } void setActiveThreads(const int);