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 9ce0ddadb6f..97f6547938a 100644 --- a/PWGJE/DataModel/JetTagging.h +++ b/PWGJE/DataModel/JetTagging.h @@ -30,6 +30,8 @@ #include #include #include +#include +#include namespace o2::aod { @@ -165,6 +167,108 @@ 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 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 { +// 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, 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 discriminant: one track-origin D_b per GnnTracks row +} // namespace jetgnnscore + +// GnnTracks has one row per evaluated (jet, track) pair. Its jet and track indices +// identify the pair independently of the ordering in JetConstituents. +// 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 +// 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_##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"); +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 +// 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(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 +// 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 +// ChargedJetGnnEdges / ChargedMCDetectorLevelJetGnnEdges. +// DECLARE_SOA_TABLE(ChargedJetGnnEdgeTags, "AOD", "CJEDGT", jetgnnscore::ScoreML); +// DECLARE_SOA_TABLE(ChargedMCDetectorLevelJetGnnEdgeTags, "AOD", "CDJEDGT", jetgnnscore::ScoreML); + } // namespace o2::aod #endif // PWGJE_DATAMODEL_JETTAGGING_H_ diff --git a/PWGJE/TableProducer/CMakeLists.txt b/PWGJE/TableProducer/CMakeLists.txt index c789aff9d49..a865a7e4b78 100644 --- a/PWGJE/TableProducer/CMakeLists.txt +++ b/PWGJE/TableProducer/CMakeLists.txt @@ -86,6 +86,16 @@ 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 + COMPONENT_NAME Analysis) + o2physics_add_dpl_workflow(estimator-rho SOURCES rhoEstimator.cxx REUSE_FROM JetDataModelPCH 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 new file mode 100644 index 00000000000..877404b8624 --- /dev/null +++ b/PWGJE/TableProducer/jetGnnInference.cxx @@ -0,0 +1,485 @@ +// 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/JetDerivedDataUtilities.h" +#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 "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 + +using namespace o2; +using namespace o2::framework; + +namespace jet_gnn_inference +{ +// Optional outputs for track-origin and future vertex prediction tables. +template +struct OptionalProduces : Produces { +}; +template <> +struct OptionalProduces { +}; +} // namespace jet_gnn_inference + +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"}; + + Configurable> binsPtMl{"binsPtMl", std::vector{5., 1000.}, "pT bin limits for ML application"}; + 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 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 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"}; + 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", "asinh", "Transformation function type for GNN"}; + + o2::analysis::MlResponseHfTagging bMlResponse; + + 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, std::vector const* trackScores = nullptr) + { + if (jet.pt() < jetpTMin) { + return; + } + 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()); + } + } + 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&) + { + 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) || + 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.init(); + + 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 (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, std::vector* trackScores = nullptr) + { + std::vector> trkFeat; + 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()}; + std::vector feat; + std::vector gnnInput; + tensorAlloc.getGNNInput(jetFeat, trkFeat, feat, gnnInput); + + 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; + } + return db; + } + + void processDummy(aod::JetCollisions::iterator const&) + { + } + PROCESS_SWITCH(JetGnnInferenceTask, processDummy, "Dummy process", true); + + void processGnnTag(typename Collisions::iterator const& collision, AnalysisJets const& jets, JetTracksExt const& tracks, CollisionContext const&) + { + 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) { + 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); + 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, processGnnFull, "Run GNN inference and fill jet tagging and track-origin prediction tables", false); +}; + +// 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-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 7ab7bc11bc6..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,10 +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 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; @@ -291,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))); @@ -312,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}}}); @@ -347,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.}}}); @@ -370,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.}}}); @@ -398,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; @@ -513,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) { @@ -560,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); @@ -582,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); @@ -594,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); @@ -613,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); //+ } } @@ -638,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*/) @@ -799,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; } @@ -827,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; } @@ -844,10 +846,7 @@ struct BJetTreeCreator { continue; } - //+ - TrackLabelMap trkLabels{{"trkVtxIndex", {}}, {"trkOrigin", {}}}; - int nVertices = jettaggingutilities::vertexClustering(collision.template mcCollision_as(), analysisJet, allTracks, MCParticles, origParticles, trkLabels, true, 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); @@ -879,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; 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..02109b6a8a9 100644 --- a/Tools/ML/model.h +++ b/Tools/ML/model.h @@ -24,6 +24,7 @@ #include +#include #include #include #include @@ -144,6 +145,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() { @@ -157,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);