diff --git a/PWGJE/Core/JetTaggingUtilities.h b/PWGJE/Core/JetTaggingUtilities.h index 754c19331dc..0d5fca0a50a 100644 --- a/PWGJE/Core/JetTaggingUtilities.h +++ b/PWGJE/Core/JetTaggingUtilities.h @@ -1155,6 +1155,65 @@ void analyzeJetTrackInfo4GNN(AnalysisJet const& analysisJet, AnyTracks const& /* } } +// Looping over the track info and putting them in the input vectors (for ParT b-jet tagging). +template +void analyzeJetTrackInfo4ParT(AnalysisJet const& analysisJet, AnyTracks const& /*allTracks*/, std::vector>& featConstit, std::vector>& vecConstit, float trackPtMin = 0.5, float trackDcaXYMax = 10.0, float trackDcaZMax = 10.0) +{ + for (const auto& constituent : analysisJet.template tracks_as()) { + + if (constituent.pt() < trackPtMin || !trackAcceptanceWithDca(constituent, trackDcaXYMax, trackDcaZMax)) { + continue; + } + + int sign = getGeoSign(analysisJet, constituent); + + float trackPt = constituent.pt(); + float trackEta = constituent.eta(); + float trackPhi = constituent.phi(); + + float ip2d = std::abs(constituent.dcaXY()) * sign; + float ipz = std::abs(constituent.dcaZ()) * sign; + float ip2dSigma = constituent.sigmadcaXY(); + float ipzSigma = constituent.sigmadcaZ(); + + float dPhi = std::atan2(std::sin(trackPhi - analysisJet.phi()), std::cos(trackPhi - analysisJet.phi())); + + featConstit.emplace_back(std::vector{ + std::log(trackPt), + std::log(trackPt / analysisJet.pt()), + jetutilities::deltaR(analysisJet, constituent), + trackEta - analysisJet.eta(), + dPhi, + std::tanh(ip2d), + std::tanh(ipz), + (ip2dSigma != 0.f) ? ip2d / ip2dSigma : 0.f, + (ipzSigma != 0.f) ? ipz / ipzSigma : 0.f, + static_cast(constituent.sign())}); + + vecConstit.emplace_back(std::vector{ + trackPt * std::cos(trackPhi), + trackPt * std::sin(trackPhi), + trackPt * std::sinh(trackEta), + trackPt * std::cosh(trackEta)}); + } +} + +// Discriminant value. Probability order is assumed to be [b, c, lf]. +template +T getDbFromProb(const std::vector& probs, double fC = 0.018) +{ + // MlResponse::getModelOutput() returns exactly nClassesMl entries, so a mismatched + // nClassesMl silently truncates the model output and would read out of bounds here. + if (probs.size() < 3) { + LOG(fatal) << "getDbFromProb expects 3 class probabilities but received " << probs.size() << ". Set the nClassesMl configurable to 3 (and give cutDirMl 3 entries to match)."; + } + T denominator = fC * probs[1] + (1. - fC) * probs[2]; + if (probs[0] <= 0. || denominator <= 0.) { + return std::numeric_limits::quiet_NaN(); + } + return std::log(probs[0] / denominator); +} + // Discriminant value for GNN b-jet tagging template T getDb(const std::vector& logits, double fC = 0.018) diff --git a/PWGJE/Core/MlResponseHfTagging.h b/PWGJE/Core/MlResponseHfTagging.h index b144d6e145a..b6cc72ed8d9 100644 --- a/PWGJE/Core/MlResponseHfTagging.h +++ b/PWGJE/Core/MlResponseHfTagging.h @@ -119,6 +119,22 @@ class MlResponseHfTagging : public MlResponse /// @return number of output nodes int getOutputNodes(int imod = 0) const { return this->mModels[imod].getNumOutputNodes(); } + /// @brief Run a model that takes several pre-built input tensors (e.g. ParT: pf_features, pf_vectors, pf_mask inputs). + /// \param input the input tensors, in the order the graph declares them + /// \param nModel index of the model + /// @return the model output, mNClasses entries long + template + std::vector getModelOutputMultiInput(T& input, int nModel = 0) + { + // evalModel returns an owning copy of the (last) output tensor of the model + std::vector output = this->mModels[nModel].template evalModel(input); + if (output.size() < this->mNClasses) { + LOG(fatal) << "Model " << this->mPaths[nModel] << " returned " << output.size() << " scores, but " << static_cast(this->mNClasses) << " classes are expected. Please check your configurables."; + } + output.resize(this->mNClasses); + return output; + } + /// Method to fill the inputs of jet, tracks and secondary vertices /// \param jet is the b-jet candidate /// \param tracks is the vector of tracks associated to the jet @@ -483,6 +499,85 @@ class GNNBjetAllocator : public TensorAllocator } }; +// TensorAllocator for the ParT b-jet tagger. +// Inputs, in graph order: pf_features (1, nFeat, nConstit), pf_vectors (1, nVec, nConstit), +// pf_mask (1, 1, nConstit). +class ParTBjetAllocator : public TensorAllocator +{ + private: + int64_t nConstit; + int64_t nFeat; + int64_t nVec; + bool padWrap; + + // Replace non-finite values in a vector, which the model cannot consume + template + static void replaceNonFinite(std::vector& vec, T value) + { + for (auto& el : vec) { // o2-linter: disable=const-ref-in-for-loop + if (std::isnan(el) || std::isinf(el)) { + el = value; + } + } + } + + public: + ParTBjetAllocator() : TensorAllocator(), nConstit(32), nFeat(10), nVec(4), padWrap(true) {} + ParTBjetAllocator(int64_t nConstit, int64_t nFeat, int64_t nVec, bool padWrap = true) + : TensorAllocator(), nConstit(nConstit), nFeat(nFeat), nVec(nVec), padWrap(padWrap) {} + ~ParTBjetAllocator() = default; + + ParTBjetAllocator& operator=(const ParTBjetAllocator& other) + { + nConstit = other.nConstit; + nFeat = other.nFeat; + nVec = other.nVec; + padWrap = other.padWrap; + return *this; + } + + // Allocate and return the ParT input tensors (std::vector). + // featData/vecData/maskData back the tensors + void getParTInput(const std::vector>& featConstit, const std::vector>& vecConstit, std::vector& featData, std::vector& vecData, std::vector& maskData, std::vector& partInput) + { + featData.assign(nFeat * nConstit, 0.f); + vecData.assign(nVec * nConstit, 0.f); + maskData.assign(nConstit, 0.f); + + int64_t nSelected = static_cast(featConstit.size()); + const int64_t nReal = (nSelected < nConstit) ? nSelected : nConstit; // jets longer than nConstit are truncated + + for (int64_t iConstit = 0; iConstit < nConstit; ++iConstit) { + int64_t iSource = iConstit; + if (iConstit < nReal) { + maskData[iConstit] = 1.f; + } else { + if (!padWrap || nReal == 0) { + continue; // zero padding, mask stays 0 + } + iSource = iConstit % nReal; // weaver 'wrap' padding: cyclic repetition, mask stays 0 + } + for (int64_t iFeat = 0; iFeat < nFeat; ++iFeat) { + featData[iFeat * nConstit + iConstit] = featConstit[iSource][iFeat]; + } + for (int64_t iVec = 0; iVec < nVec; ++iVec) { + vecData[iVec * nConstit + iConstit] = vecConstit[iSource][iVec]; + } + } + + replaceNonFinite(featData, 0.f); + replaceNonFinite(vecData, 0.f); + + std::vector featShape{1, nFeat, nConstit}; + std::vector vecShape{1, nVec, nConstit}; + std::vector maskShape{1, 1, nConstit}; + + partInput.emplace_back(createTensor(featData, featShape)); + partInput.emplace_back(createTensor(vecData, vecShape)); + partInput.emplace_back(createTensor(maskData, maskShape)); + } +}; + } // namespace o2::analysis #undef FILL_MAP_BJET diff --git a/PWGJE/TableProducer/jetTaggerHF.cxx b/PWGJE/TableProducer/jetTaggerHF.cxx index bd8625b676c..f2daa6a36e9 100644 --- a/PWGJE/TableProducer/jetTaggerHF.cxx +++ b/PWGJE/TableProducer/jetTaggerHF.cxx @@ -145,6 +145,17 @@ struct JetTaggerHFTask { ConfigurableAxis binTrackProbability{"binTrackProbability", {100, 0.f, 1.f}, ""}; ConfigurableAxis binJetFlavour{"binJetFlavour", {6, -0.5, 5.5}, ""}; + // ParT (Particle Transformer) tagger. Grouped in a ConfigurableGroup because + // StructToTuple.h supports at most 99 task members and this task is near the limit. + struct ParTCfg : ConfigurableGroup { + std::string prefix = "part"; + Configurable nConstit{"nConstit", 32, "fixed number of constituents fed to the ParT model, must match training"}; + Configurable nFeat{"nFeat", 10, "number of pf_features channels"}; + Configurable nVec{"nVec", 4, "number of pf_vectors channels"}; + Configurable padWrap{"padWrap", true, "pad short jets by cyclic repetition (weaver 'wrap' mode); set false for zero padding"}; + Configurable minTracks{"minTracks", 2, "minimum number of selected constituents; training applied jet_nTracks >= 2"}; + } parTCfg; + o2::analysis::MlResponseHfTagging bMlResponse; o2::ccdb::CcdbApi ccdbApi; @@ -277,7 +288,7 @@ struct JetTaggerHFTask { } } } - if (doprocessAlgorithmGNN || doprocessAlgorithmGNNwExtra) { + if (doprocessAlgorithmGNN || doprocessAlgorithmGNNwExtra || doprocessAlgorithmParT) { if (jet.pt() >= jetpTMin) { float dbRange; if (scoreML[jet.globalIndex()] < dbMin) { @@ -525,9 +536,24 @@ struct JetTaggerHFTask { bMlResponse.init(); } + if (doprocessAlgorithmParT) { + bMlResponse.configure(binsPtMl, cutsMl, cutDirMl, nClassesMl); + if (loadModelsFromCCDB) { + ccdbApi.init(ccdbUrl); + bMlResponse.setModelPathsCCDB(onnxFileNames, ccdbApi, modelPathsCCDB, timestampCCDB); + } else { + bMlResponse.setModelPathsLocal(onnxFileNames); + } + bMlResponse.init(); + + LOGF(info, "ParT: loaded %s with %d output node(s); %d constituents, %s padding", onnxFileNames.value[0].c_str(), bMlResponse.getOutputNodes(), static_cast(parTCfg.nConstit.value), parTCfg.padWrap.value ? "wrap" : "zero"); + } + if (doprocessAlgorithmGNN || doprocessAlgorithmGNNwExtra) { tensorAlloc = o2::analysis::GNNBjetAllocator(nJetFeat.value, nTrkFeat.value, nClassesMl.value, nTrkOrigin.value, transformFeatureJetMean.value, transformFeatureJetStdev.value, transformFeatureTrkMean.value, transformFeatureTrkStdev.value, nJetConst, tfFuncTypeGNN.value); + } + if (doprocessAlgorithmGNN || doprocessAlgorithmGNNwExtra || doprocessAlgorithmParT) { 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"); @@ -669,6 +695,42 @@ struct JetTaggerHFTask { } } + template + void analyzeJetAlgorithmParT(AnyJets const& jets, AnyTracks const& tracks) + { + o2::analysis::ParTBjetAllocator parTTensorAlloc(parTCfg.nConstit.value, parTCfg.nFeat.value, parTCfg.nVec.value, parTCfg.padWrap.value); + + for (const auto& jet : jets) { + std::vector> featConstit; + std::vector> vecConstit; + jettaggingutilities::analyzeJetTrackInfo4ParT(jet, tracks, featConstit, vecConstit, trackPtMin, trackDcaXYMax, trackDcaZMax); + + // Training applied a jet_nTracks >= 2 selection, so jets below it are out of distribution. + if (static_cast(featConstit.size()) < parTCfg.minTracks) { + scoreML[jet.globalIndex()] = -999.; + LOGF(debug, "doprocessAlgorithmParT, too few selected constituents (%d)", jet.globalIndex()); + continue; + } + + std::vector featData; + std::vector vecData; + std::vector maskData; + std::vector parTInput; + parTTensorAlloc.getParTInput(featConstit, vecConstit, featData, vecData, maskData, parTInput); + + // The exported graph applies the softmax itself, so the output is already + // [p_b, p_c, p_udsg] and must not be passed through getDb(), which would softmax again. + auto modelOutput = bMlResponse.getModelOutputMultiInput(parTInput, 0); + float db = jettaggingutilities::getDbFromProb(modelOutput, fC); + if (!std::isnan(db)) { + scoreML[jet.globalIndex()] = db; + } else { + scoreML[jet.globalIndex()] = 999.; + LOGF(debug, "doprocessAlgorithmParT, Db is NaN (%d)", jet.globalIndex()); + } + } + } + void processDummy(aod::JetCollisions const&) { } @@ -725,6 +787,12 @@ struct JetTaggerHFTask { } PROCESS_SWITCH(JetTaggerHFTask, processAlgorithmGNN, "Fill GNN evaluation score (D_b) for charged jets", false); + void processAlgorithmParT(JetTable const& jets, JetTracksExt const& jtracks) + { + analyzeJetAlgorithmParT(jets, jtracks); + } + PROCESS_SWITCH(JetTaggerHFTask, processAlgorithmParT, "Fill ParT evaluation score (D_b) for charged jets", false); + void processFillTables(std::conditional_t, JetTable>::iterator const& jet, JetTracksExt const& tracks) { fillTables(jet, tracks);