[PWGJE] Add ParT b-jet tagger to jetTaggerHF

Author: ahmadtabikh

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 <typename AnalysisJet, typename AnyTracks>
+void analyzeJetTrackInfo4ParT(AnalysisJet const& analysisJet, AnyTracks const& /*allTracks*/, std::vector<std::vector<float>>& featConstit, std::vector<std::vector<float>>& vecConstit, float trackPtMin = 0.5, float trackDcaXYMax = 10.0, float trackDcaZMax = 10.0)
+{
+  for (const auto& constituent : analysisJet.template tracks_as<AnyTracks>()) {
+
+    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<float>{
+      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<float>(constituent.sign())});
+
+    vecConstit.emplace_back(std::vector<float>{
+      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 <typename T>
+T getDbFromProb(const std::vector<T>& 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<T>::quiet_NaN();
+  }
+  return std::log(probs[0] / denominator);
+}
+
 // Discriminant value for GNN b-jet tagging
 template <typename T>
 T getDb(const std::vector<T>& logits, double fC = 0.018)

PWGJE/Core/MlResponseHfTagging.h

@@ -119,6 +119,22 @@ class MlResponseHfTagging : public MlResponse<TypeOutputScore>
   /// @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 <typename T>
+  std::vector<TypeOutputScore> getModelOutputMultiInput(T& input, int nModel = 0)
+  {
+    // evalModel returns an owning copy of the (last) output tensor of the model
+    std::vector<TypeOutputScore> output = this->mModels[nModel].template evalModel<TypeOutputScore>(input);
+    if (output.size() < this->mNClasses) {
+      LOG(fatal) << "Model " << this->mPaths[nModel] << " returned " << output.size() << " scores, but " << static_cast<int>(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 <typename T>
+  static void replaceNonFinite(std::vector<T>& 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<Ort::Value>).
+  // featData/vecData/maskData back the tensors
+  void getParTInput(const std::vector<std::vector<float>>& featConstit, const std::vector<std::vector<float>>& vecConstit, std::vector<float>& featData, std::vector<float>& vecData, std::vector<float>& maskData, std::vector<Ort::Value>& partInput)
+  {
+    featData.assign(nFeat * nConstit, 0.f);
+    vecData.assign(nVec * nConstit, 0.f);
+    maskData.assign(nConstit, 0.f);
+
+    int64_t nSelected = static_cast<int64_t>(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<int64_t> featShape{1, nFeat, nConstit};
+    std::vector<int64_t> vecShape{1, nVec, nConstit};
+    std::vector<int64_t> 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

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<int64_t> nConstit{"nConstit", 32, "fixed number of constituents fed to the ParT model, must match training"};
+    Configurable<int64_t> nFeat{"nFeat", 10, "number of pf_features channels"};
+    Configurable<int64_t> nVec{"nVec", 4, "number of pf_vectors channels"};
+    Configurable<bool> padWrap{"padWrap", true, "pad short jets by cyclic repetition (weaver 'wrap' mode); set false for zero padding"};
+    Configurable<int> minTracks{"minTracks", 2, "minimum number of selected constituents; training applied jet_nTracks >= 2"};
+  } parTCfg;
+
   o2::analysis::MlResponseHfTagging<float> 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<int>(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<TH2>(HIST("h2_count_db"));
       h2CountDb->GetXaxis()->SetBinLabel(1, "b-jet");
@@ -669,6 +695,42 @@ struct JetTaggerHFTask {
     }
   }
 
+  template <typename AnyJets, typename AnyTracks>
+  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<std::vector<float>> featConstit;
+      std::vector<std::vector<float>> 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<int>(featConstit.size()) < parTCfg.minTracks) {
+        scoreML[jet.globalIndex()] = -999.;
+        LOGF(debug, "doprocessAlgorithmParT, too few selected constituents (%d)", jet.globalIndex());
+        continue;
+      }
+
+      std::vector<float> featData;
+      std::vector<float> vecData;
+      std::vector<float> maskData;
+      std::vector<Ort::Value> 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<isMCD, soa::Join<JetTable, aod::ChargedMCDetectorLevelJetFlavourDef>, JetTable>::iterator const& jet, JetTracksExt const& tracks)
   {
     fillTables<isMCD>(jet, tracks);