[PWGJE] Add cut selection for IP and mass distribution of sv tagging

PWGJE/Core/JetTaggingUtilities.h

@@ -304,7 +304,7 @@ int16_t getJetFlavor(AnyJet const& jet, AllMCParticles const& mcparticles)
 
       if (dR < jet.r() / 100.f) {
         if (TMath::Abs(pdgcode) == 5) {
-          return 2; // Beauty jet
+          return JetTaggingSpecies::beauty; // Beauty jet
         } else {
           if (count > arraySize - 1)
             return 0;
@@ -317,10 +317,41 @@ int16_t getJetFlavor(AnyJet const& jet, AllMCParticles const& mcparticles)
 
   for (int ij = 0; ij < count; ij++) {
     if (TMath::Abs(countpartcode[ij]) == 4)
-      return 1; // Charm jet
+      return JetTaggingSpecies::charm; // Charm jet
   }
 
-  return 0; // Light flavor jet
+  return JetTaggingSpecies::lightflavour; // Light flavor jet
+}
+
+/**
+ * return acceptance of track about DCA xy and z due to cut for QualityTracks
+ */
+template <typename T>
+bool trackAcceptanceWithDca(T const& track, float trackDcaXYMax, float trackDcaZMax)
+{
+  if (std::abs(track.dcaXY()) > trackDcaXYMax)
+    return false;
+  if (std::abs(track.dcaZ()) > trackDcaZMax)
+    return false;
+  return true;
+}
+
+/**
+ * retrun acceptance of prong about chi2 and error of decay length due to cut for high quality secondary vertex
+ */
+template <typename T>
+bool prongAcceptance(T const& prong, float prongChi2PCAMax, float prongsigmaLxyMax, bool doXYZ)
+{
+  if (prong.chi2PCA() > prongChi2PCAMax)
+    return false;
+  if (!doXYZ) {
+    if (prong.errorDecayLengthXY() > prongsigmaLxyMax)
+      return false;
+  } else {
+    if (prong.errorDecayLength() > prongsigmaLxyMax)
+      return false;
+  }
+  return true;
 }
 
 /**
@@ -350,10 +381,12 @@ int getGeoSign(T const& collision, U const& jet, V const& track)
  * in a vector in descending order.
  */
 template <typename T, typename U, typename V, typename W, typename Vec = std::vector<float>>
-void orderForIPJetTracks(T const& collision, U const& jet, V const& /*jtracks*/, W const& /*tracks*/, Vec& vecSignImpSig)
+void orderForIPJetTracks(T const& collision, U const& jet, V const& /*jtracks*/, W const& /*tracks*/, float const& trackDcaXYMax, float const& trackDcaZMax, Vec& vecSignImpSig)
 {
   for (auto& jtrack : jet.template tracks_as<V>()) {
     auto track = jtrack.template track_as<W>();
+    if (!trackAcceptanceWithDca(track, trackDcaXYMax, trackDcaZMax))
+      continue;
     auto geoSign = getGeoSign(collision, jet, track);
     auto varSignImpXYSig = geoSign * TMath::Abs(track.dcaXY()) / TMath::Sqrt(track.sigmaDcaXY2());
     vecSignImpSig.push_back(varSignImpXYSig);
@@ -364,14 +397,14 @@ void orderForIPJetTracks(T const& collision, U const& jet, V const& /*jtracks*/,
 /**
  * Checks if a jet is greater than the given tagging working point based on the signed impact parameter significances
  */
-template <typename T, typename U, typename V, typename W, typename X, typename Y>
-bool isGreaterThanTaggingPoint(T const& collision, U const& jet, V const& jtracks, W const& tracks, X const& taggingPoint = 1.0, Y const& cnt = 1)
+template <typename T, typename U, typename V, typename W>
+bool isGreaterThanTaggingPoint(T const& collision, U const& jet, V const& jtracks, W const& tracks, float const& trackDcaXYMax, float const& trackDcaZMax, float const& taggingPoint = 1.0, int const& cnt = 1)
 {
   if (cnt == 0) {
     return true; // untagged
   }
   std::vector<float> vecSignImpSig;
-  orderForIPJetTracks(collision, jet, jtracks, tracks, vecSignImpSig);
+  orderForIPJetTracks(collision, jet, jtracks, tracks, trackDcaXYMax, trackDcaZMax, vecSignImpSig);
   if (vecSignImpSig.size() > static_cast<std::vector<float>::size_type>(cnt) - 1) {
     for (int i = 0; i < cnt; i++) {
       if (vecSignImpSig[i] < taggingPoint) { // tagger point set
@@ -448,16 +481,18 @@ float getTrackProbability(T const& fResoFuncjet, U const& track, const float& mi
  *         geometric sign.
  */
 template <typename T, typename U, typename V, typename W, typename X>
-float getJetProbability(T const& fResoFuncjet, U const& collision, V const& jet, W const& jtracks, X const& tracks, const int& cnt, const float& tagPoint = 1.0, const float& minSignImpXYSig = -10)
+float getJetProbability(T const& fResoFuncjet, U const& collision, V const& jet, W const& jtracks, X const& tracks, float const& trackDcaXYMax, float const& trackDcaZMax, const int& cnt, const float& tagPoint = 1.0, const float& minSignImpXYSig = -10)
 {
-  if (!(isGreaterThanTaggingPoint(collision, jet, jtracks, tracks, tagPoint, cnt)))
+  if (!(isGreaterThanTaggingPoint(collision, jet, jtracks, tracks, trackDcaXYMax, trackDcaZMax, tagPoint, cnt)))
     return -1;
 
   std::vector<float> jetTracksPt;
   float trackjetProb = 1.;
 
   for (auto& jtrack : jet.template tracks_as<W>()) {
     auto track = jtrack.template track_as<X>();
+    if (!trackAcceptanceWithDca(track, trackDcaXYMax, trackDcaZMax))
+      continue;
 
     float probTrack = getTrackProbability(fResoFuncjet, track, minSignImpXYSig);
 
@@ -481,6 +516,177 @@ float getJetProbability(T const& fResoFuncjet, U const& collision, V const& jet,
   return JP;
 }
 
+// For secaondy vertex method utilites
+class bjetCandSV
+{
+ public:
+  bjetCandSV() = default;
+
+  bjetCandSV(float xpv, float ypv, float zpv, float xsv, float ysv, float zsv,
+             float pxVal, float pyVal, float pzVal, float eVal, float mVal, float chi2Val,
+             float errDecayLength, float errDecayLengthXY,
+             float rSecVertex, float ptVal, float pVal,
+             std::array<float, 3> pVec, float etaVal, float phiVal,
+             float yVal, float decayLen, float decayLenXY,
+             float decayLenNorm, float decayLenXYNorm,
+             float cpaVal, float impParXY)
+    : m_xPVertex(xpv), m_yPVertex(ypv), m_zPVertex(zpv), m_xSecondaryVertex(xsv), m_ySecondaryVertex(ysv), m_zSecondaryVertex(zsv), m_px(pxVal), m_py(pyVal), m_pz(pzVal), m_e(eVal), m_m(mVal), m_chi2PCA(chi2Val), m_errorDecayLength(errDecayLength), m_errorDecayLengthXY(errDecayLengthXY), m_rSecondaryVertex(rSecVertex), m_pt(ptVal), m_p(pVal), m_pVector(pVec), m_eta(etaVal), m_phi(phiVal), m_y(yVal), m_decayLength(decayLen), m_decayLengthXY(decayLenXY), m_decayLengthNormalised(decayLenNorm), m_decayLengthXYNormalised(decayLenXYNorm), m_cpa(cpaVal), m_impactParameterXY(impParXY)
+  {
+  }
+
+  float xPVertex() const { return m_xPVertex; }
+  float yPVertex() const { return m_yPVertex; }
+  float zPVertex() const { return m_zPVertex; }
+
+  float xSecondaryVertex() const { return m_xSecondaryVertex; }
+  float ySecondaryVertex() const { return m_ySecondaryVertex; }
+  float zSecondaryVertex() const { return m_zSecondaryVertex; }
+
+  float px() const { return m_px; }
+  float py() const { return m_py; }
+  float pz() const { return m_pz; }
+  float e() const { return m_e; }
+  float m() const { return m_m; }
+  float chi2PCA() const { return m_chi2PCA; }
+
+  float errorDecayLength() const { return m_errorDecayLength; }
+  float errorDecayLengthXY() const { return m_errorDecayLengthXY; }
+
+  float rSecondaryVertex() const { return m_rSecondaryVertex; }
+  float pt() const { return m_pt; }
+  float p() const { return m_p; }
+
+  std::array<float, 3> pVector() const { return m_pVector; }
+
+  float eta() const { return m_eta; }
+  float phi() const { return m_phi; }
+  float y() const { return m_y; }
+
+  float decayLength() const { return m_decayLength; }
+  float decayLengthXY() const { return m_decayLengthXY; }
+  float decayLengthNormalised() const { return m_decayLengthNormalised; }
+  float decayLengthXYNormalised() const { return m_decayLengthXYNormalised; }
+
+  float cpa() const { return m_cpa; }
+  float impactParameterXY() const { return m_impactParameterXY; }
+
+ private:
+  float m_xPVertex, m_yPVertex, m_zPVertex;
+  float m_xSecondaryVertex, m_ySecondaryVertex, m_zSecondaryVertex;
+  float m_px, m_py, m_pz, m_e, m_m, m_chi2PCA;
+  float m_errorDecayLength, m_errorDecayLengthXY;
+  float m_rSecondaryVertex, m_pt, m_p;
+  std::array<float, 3> m_pVector;
+  float m_eta, m_phi, m_y;
+  float m_decayLength, m_decayLengthXY, m_decayLengthNormalised, m_decayLengthXYNormalised;
+  float m_cpa, m_impactParameterXY;
+};
+
+template <typename ProngType, typename JetType>
+bjetCandSV jetFromProngMaxDecayLength(const JetType& jet, float const& prongChi2PCAMax, float const& prongsigmaLxyMax, const bool& doXYZ = false)
+{
+  float xPVertex = 0.0f;
+  float yPVertex = 0.0f;
+  float zPVertex = 0.0f;
+  float xSecondaryVertex = 0.0f;
+  float ySecondaryVertex = 0.0f;
+  float zSecondaryVertex = 0.0f;
+  float px = 0.0f;
+  float py = 0.0f;
+  float pz = 0.0f;
+  float e = 0.0f;
+  float m = 0.0f;
+  float chi2PCA = 0.0f;
+  float errorDecayLength = 0.0f;
+  float errorDecayLengthXY = 0.0f;
+
+  float rSecondaryVertex = 0.0f;
+  float pt = 0.0f;
+  float p = 0.0f;
+  std::array<float, 3> pVector = {0.0f, 0.0f, 0.0f};
+  float eta = 0.0f;
+  float phi = 0.0f;
+  float y = 0.0f;
+  float decayLength = 0.0f;
+  float decayLengthXY = 0.0f;
+  float decayLengthNormalised = 0.0f;
+  float decayLengthXYNormalised = 0.0f;
+  float cpa = 0.0f;
+  float impactParameterXY = 0.0f;
+
+  float maxSxy = -1.0f;
+
+  for (const auto& prong : jet.template secondaryVertices_as<ProngType>()) {
+    float Sxy = -1.;
+    if (!doXYZ) {
+      Sxy = prong.decayLengthXY() / prong.errorDecayLengthXY();
+    } else {
+      Sxy = prong.decayLength() / prong.errorDecayLength();
+    }
+    if (!prongAcceptance(prong, prongChi2PCAMax, prongsigmaLxyMax, doXYZ))
+      continue;
+
+    if (maxSxy < Sxy) {
+      maxSxy = Sxy;
+
+      xPVertex = prong.xPVertex();
+      yPVertex = prong.yPVertex();
+      zPVertex = prong.zPVertex();
+      xSecondaryVertex = prong.xSecondaryVertex();
+      ySecondaryVertex = prong.ySecondaryVertex();
+      zSecondaryVertex = prong.zSecondaryVertex();
+      px = prong.px();
+      py = prong.py();
+      pz = prong.pz();
+      e = prong.e();
+      m = prong.m();
+      chi2PCA = prong.chi2PCA();
+      errorDecayLength = prong.errorDecayLength();
+      errorDecayLengthXY = prong.errorDecayLengthXY();
+      rSecondaryVertex = prong.rSecondaryVertex();
+      pt = prong.pt();
+      p = prong.p();
+      pVector = prong.pVector();
+      eta = prong.eta();
+      phi = prong.phi();
+      y = prong.y();
+      decayLength = prong.decayLength();
+      decayLengthXY = prong.decayLengthXY();
+      decayLengthNormalised = prong.decayLengthNormalised();
+      decayLengthXYNormalised = prong.decayLengthXYNormalised();
+      cpa = prong.cpa();
+      impactParameterXY = prong.impactParameterXY();
+    }
+  }
+
+  return bjetCandSV(
+    xPVertex, yPVertex, zPVertex,
+    xSecondaryVertex, ySecondaryVertex, zSecondaryVertex,
+    px, py, pz, e, m, chi2PCA,
+    errorDecayLength, errorDecayLengthXY,
+    rSecondaryVertex, pt, p,
+    pVector, eta, phi,
+    y, decayLength, decayLengthXY,
+    decayLengthNormalised, decayLengthXYNormalised,
+    cpa, impactParameterXY);
+}
+
+template <typename T, typename U>
+bool isTaggedJetSV(T const jet, U const& /*prongs*/, float const& prongChi2PCAMax, float const& prongsigmaLxyMax, float const& doXYZ = false, float const& tagPointForSV = 15.)
+{
+  auto bjetCand = jetFromProngMaxDecayLength<U>(jet, prongChi2PCAMax, prongsigmaLxyMax, doXYZ);
+  if (!doXYZ) {
+    auto maxSxy = bjetCand.decayLengthXY() / bjetCand.errorDecayLengthXY();
+    if (maxSxy < tagPointForSV)
+      return false;
+  } else {
+    auto maxSxyz = bjetCand.decayLength() / bjetCand.errorDecayLength();
+    if (maxSxyz < tagPointForSV)
+      return false;
+  }
+  return true;
+}
+
 }; // namespace jettaggingutilities
 
 #endif // PWGJE_CORE_JETTAGGINGUTILITIES_H_

PWGJE/DataModel/JetTagging.h

@@ -123,10 +123,10 @@ JETSV_TABLES_DEF(Charged, SecondaryVertex2Prong, "2PRONG");
   {                                                             \
     DECLARE_SOA_COLUMN(Origin, origin, int);                    \
     DECLARE_SOA_COLUMN(JetProb, jetProb, std::vector<float>);   \
-    DECLARE_SOA_COLUMN(Algorithm2, algorithm2, int);            \
-    DECLARE_SOA_COLUMN(Algorithm3, algorithm3, int);            \
+    DECLARE_SOA_COLUMN(FlagtaggedjetIP, flagtaggedjetIP, bool); \
+    DECLARE_SOA_COLUMN(FlagtaggedjetSV, flagtaggedjetSV, bool); \
   }                                                             \
-  DECLARE_SOA_TABLE(_jet_type_##Tags, "AOD", _description_ "Tags", _name_##tagging::Origin, _name_##tagging::JetProb, _name_##tagging::Algorithm2, _name_##tagging::Algorithm3);
+  DECLARE_SOA_TABLE(_jet_type_##Tags, "AOD", _description_ "Tags", _name_##tagging::Origin, _name_##tagging::JetProb, _name_##tagging::FlagtaggedjetIP, _name_##tagging::FlagtaggedjetSV);
 
 #define JETTAGGING_TABLES_DEF(_jet_type_, _description_)                                                    \
   JETTAGGING_TABLE_DEF(_jet_type_##Jet, _jet_type_##jet, _description_)                                     \