[PWGHF] KFParticle based 3 prong decay reconstruction development

PWGHF/Core/SelectorCuts.h

@@ -465,6 +465,7 @@ namespace hf_cuts_lc_to_p_k_pi
 {
 static constexpr int NBinsPt = 10;
 static constexpr int NCutVars = 11;
+static constexpr int NCutKfVars = 12;
 // default values for the pT bin edges (can be used to configure histogram axis)
 // offset by 1 from the bin numbers in cuts array
 constexpr double BinsPt[NBinsPt + 1] = {
@@ -493,6 +494,19 @@ constexpr double Cuts[NBinsPt][NCutVars] = {{0.4, 0.4, 0.4, 0.4, 0., 0.005, 0.,
                                             {0.4, 0.4, 0.4, 0.4, 0., 0.005, 0., 0., 0., 1e+10, -1.},  /* 12 < pT < 24 */
                                             {0.4, 0.4, 0.4, 0.4, 0., 0.005, 0., 0., 0., 1e+10, -1.}}; /* 24 < pT < 36 */
 
+// default value for the cuts                    Chi2Prim      Chi2Geo       DCA, cm     Chi2Geo  Chi2Topo
+//                                               P   K  Pi  KPi PPi  PK  KPi    PPi   PK   ↓  LdL  ↓
+constexpr double CutsKf[NBinsPt][NCutKfVars] = {{3., 3., 3., 3., 3., 3., 0.01, 0.01, 0.01, 3., 5., 5.},  /* 0  < pT < 1  */
+                                                {3., 3., 3., 3., 3., 3., 0.01, 0.01, 0.01, 3., 5., 5.},  /* 1  < pT < 2  */
+                                                {3., 3., 3., 3., 3., 3., 0.01, 0.01, 0.01, 3., 5., 5.},  /* 2  < pT < 3  */
+                                                {3., 3., 3., 3., 3., 3., 0.01, 0.01, 0.01, 3., 5., 5.},  /* 3  < pT < 4  */
+                                                {3., 3., 3., 3., 3., 3., 0.01, 0.01, 0.01, 3., 5., 5.},  /* 4  < pT < 5  */
+                                                {3., 3., 3., 3., 3., 3., 0.01, 0.01, 0.01, 3., 5., 5.},  /* 5  < pT < 6  */
+                                                {3., 3., 3., 3., 3., 3., 0.01, 0.01, 0.01, 3., 5., 5.},  /* 6  < pT < 8  */
+                                                {3., 3., 3., 3., 3., 3., 0.01, 0.01, 0.01, 3., 5., 5.},  /* 8  < pT < 12 */
+                                                {3., 3., 3., 3., 3., 3., 0.01, 0.01, 0.01, 3., 5., 5.},  /* 12 < pT < 24 */
+                                                {3., 3., 3., 3., 3., 3., 0.01, 0.01, 0.01, 3., 5., 5.}}; /* 24 < pT < 36 */
+
 // row labels
 static const std::vector<std::string> labelsPt = {
   "pT bin 0",
@@ -508,6 +522,7 @@ static const std::vector<std::string> labelsPt = {
 
 // column labels
 static const std::vector<std::string> labelsCutVar = {"m", "pT p", "pT K", "pT Pi", "Chi2PCA", "decay length", "cos pointing angle", "decLengthXY", "normDecLXY", "impParXY", "mass (Kpi)"};
+static const std::vector<std::string> labelsCutKfVar = {"kfChi2PrimPr", "kfChi2PrimKa", "kfChi2PrimPi", "kfChi2GeoKaPi", "kfChi2GeoPrPi", "kfChi2GeoPrKa", "kfDcaKaPi", "kfDcaPrPi", "kfDcaPrKa", "kfChi2Geo", "kfLdL", "kfChi2Topo"};
 } // namespace hf_cuts_lc_to_p_k_pi
 
 namespace hf_cuts_lc_to_k0s_p

PWGHF/TableProducer/candidateCreator3Prong.cxx

@@ -85,6 +85,9 @@
   Configurable<bool> createDs{"createDs", false, "enable Ds+/- candidate creation"};
   Configurable<bool> createLc{"createLc", false, "enable Lc+/- candidate creation"};
   Configurable<bool> createXic{"createXic", false, "enable Xic+/- candidate creation"};
+  // KF
+  Configurable<bool> applyTopoConstraint{"applyTopoConstraint", false, "apply origin from PV hypothesis for created candidate, works only in KF mode"};
+  Configurable<bool> applyInvMassConstraint{"applyInvMassConstraint", false, "apply particle type hypothesis to recalculate created candidate's momentum, works only in KF mode"};
 
   HfEventSelection hfEvSel;        // event selection and monitoring
   o2::vertexing::DCAFitterN<3> df; // 3-prong vertex fitter
@@ -149,6 +152,10 @@
       LOGP(fatal, "At least one particle specie should be enabled for the creation.");
     }
 
+    if (createLc && createXic && applyInvMassConstraint) {
+      LOGP(fatal, "Unable to apply invariant mass constraint due to ambiguity of mass hypothesis: only one of Lc and Xic can be reconstructed.");
+    }
+
     // histograms
     registry.add("hMass3PKPi", "3-prong candidates;inv. mass (pK#pi) (GeV/#it{c}^{2});entries", {HistType::kTH1D, {{1200, 1.8, 3.0}}});
     registry.add("hMass3PiKP", "3-prong candidates;inv. mass (#pi Kp) (GeV/#it{c}^{2});entries", {HistType::kTH1D, {{1200, 1.8, 3.0}}});
@@ -487,6 +494,15 @@
       kfCandPiKK.SetConstructMethod(2);
       kfCandPiKK.Construct(kfDaughtersPiKK, 3);
 
+      const float chi2topo = kfCalculateChi2ToPrimaryVertex(kfCandPKPi, kfpV);
+
+      if (applyTopoConstraint) { // constraints applied after chi2topo getter - to preserve unbiased value of chi2topo
+        for (const auto& kfCand : std::array<KFParticle*, 5>{&kfCandPKPi, &kfCandPiKP, &kfCandPiKPi, &kfCandKKPi, &kfCandPiKK}) {
+          kfCand->SetProductionVertex(kfpV);
+          kfCand->TransportToDecayVertex();
+        }
+      }
+
       KFParticle kfPairKPi;
       const KFParticle* kfDaughtersKPi[3] = {&kfSecondKaon, &kfThirdPion};
       kfPairKPi.SetConstructMethod(2);
@@ -505,8 +521,15 @@
       const float massKPi = kfPairKPi.GetMass();
       const float massPiK = kfPairPiK.GetMass();
 
+      if (applyInvMassConstraint) { // constraints applied after minv getters - to preserve unbiased values of minv
+        kfCandPKPi.SetNonlinearMassConstraint(createLc ? MassLambdaCPlus : MassXiCPlus);
+        kfCandPiKP.SetNonlinearMassConstraint(createLc ? MassLambdaCPlus : MassXiCPlus);
+        kfCandPiKPi.SetNonlinearMassConstraint(MassDPlus);
+        kfCandKKPi.SetNonlinearMassConstraint(MassDS);
+        kfCandPiKK.SetNonlinearMassConstraint(MassDS);
+      }
+
       const float chi2geo = kfCandPKPi.Chi2() / kfCandPKPi.NDF();
-      const float chi2topo = kfCalculateChi2ToPrimaryVertex(kfCandPKPi, kfpV);
       const std::pair<float, float> ldl = kfCalculateLdL(kfCandPKPi, kfpV);
 
       std::array<float, 3> pProng0 = kfCalculateProngMomentumInSecondaryVertex(kfFirstProton, kfCandPiKP);
@@ -924,7 +947,7 @@
             swapping = int8_t(std::abs(arrayDaughters[0].mcParticle().pdgCode()) == kPiPlus);
           }
           RecoDecay::getDaughters(mcParticles.rawIteratorAt(indexRec), &arrDaughIndex, std::array{0}, 1);
           if (arrDaughIndex.size() == 2) {
             for (auto iProng = 0u; iProng < arrDaughIndex.size(); ++iProng) {
               auto daughI = mcParticles.rawIteratorAt(arrDaughIndex[iProng]);
               arrPDGDaugh[iProng] = std::abs(daughI.pdgCode());
@@ -970,7 +993,7 @@
             swapping = int8_t(std::abs(arrayDaughters[0].mcParticle().pdgCode()) == kPiPlus);
           }
           RecoDecay::getDaughters(mcParticles.rawIteratorAt(indexRec), &arrDaughIndex, std::array{0}, 1);
           if (arrDaughIndex.size() == 2) {
             for (auto iProng = 0u; iProng < arrDaughIndex.size(); ++iProng) {
               auto daughI = mcParticles.rawIteratorAt(arrDaughIndex[iProng]);
               arrPDGDaugh[iProng] = std::abs(daughI.pdgCode());

PWGHF/TableProducer/candidateSelectorLc.cxx

@@ -73,6 +73,9 @@ struct HfCandidateSelectorLc {
   // topological cuts
   Configurable<std::vector<double>> binsPt{"binsPt", std::vector<double>{hf_cuts_lc_to_p_k_pi::vecBinsPt}, "pT bin limits"};
   Configurable<LabeledArray<double>> cuts{"cuts", {hf_cuts_lc_to_p_k_pi::Cuts[0], hf_cuts_lc_to_p_k_pi::NBinsPt, hf_cuts_lc_to_p_k_pi::NCutVars, hf_cuts_lc_to_p_k_pi::labelsPt, hf_cuts_lc_to_p_k_pi::labelsCutVar}, "Lc candidate selection per pT bin"};
+  Configurable<LabeledArray<double>> kfCuts{"kfCuts", {hf_cuts_lc_to_p_k_pi::CutsKf[0], hf_cuts_lc_to_p_k_pi::NBinsPt, hf_cuts_lc_to_p_k_pi::NCutKfVars, hf_cuts_lc_to_p_k_pi::labelsPt, hf_cuts_lc_to_p_k_pi::labelsCutKfVar}, "Lc candidate selection per pT bin with KF-associated variables"};
+  Configurable<bool> applyNonKfCuts{"applyNonKfCuts", true, "Whether to apply non-KF cuts when running in KF mode. In DCAFitter mode this field is not effective"};
+  Configurable<bool> applyKfCuts{"applyKfCuts", true, "Whether to apply KF cuts when running in KF mode. In DCAFitter mode this field is not effective"};
   // QA switch
   Configurable<bool> activateQA{"activateQA", false, "Flag to enable QA histogram"};
   // ML inference
@@ -175,7 +178,7 @@ struct HfCandidateSelectorLc {
   /// Conjugate-independent topological cuts
   /// \param candidate is candidate
   /// \return true if candidate passes all cuts
-  template <typename T>
+  template <aod::hf_cand::VertexerType reconstructionType, typename T>
   bool selectionTopol(const T& candidate)
   {
     auto candpT = candidate.pt();
@@ -190,33 +193,54 @@ struct HfCandidateSelectorLc {
       return false;
     }
 
-    // cosine of pointing angle
-    if (candidate.cpa() <= cuts->get(pTBin, "cos pointing angle")) {
-      return false;
-    }
+    if (reconstructionType == aod::hf_cand::VertexerType::DCAFitter || (reconstructionType == aod::hf_cand::VertexerType::KfParticle && applyNonKfCuts)) {
+      // cosine of pointing angle
+      if (candidate.cpa() <= cuts->get(pTBin, "cos pointing angle")) {
+        return false;
+      }
 
-    // candidate chi2PCA
-    if (candidate.chi2PCA() > cuts->get(pTBin, "Chi2PCA")) {
-      return false;
-    }
+      // candidate chi2PCA
+      if (candidate.chi2PCA() > cuts->get(pTBin, "Chi2PCA")) {
+        return false;
+      }
 
-    if (candidate.decayLength() <= cuts->get(pTBin, "decay length")) {
-      return false;
-    }
+      if (candidate.decayLength() <= cuts->get(pTBin, "decay length")) {
+        return false;
+      }
 
-    // candidate decay length XY
-    if (candidate.decayLengthXY() <= cuts->get(pTBin, "decLengthXY")) {
-      return false;
-    }
+      // candidate decay length XY
+      if (candidate.decayLengthXY() <= cuts->get(pTBin, "decLengthXY")) {
+        return false;
+      }
 
-    // candidate normalized decay length XY
-    if (candidate.decayLengthXYNormalised() < cuts->get(pTBin, "normDecLXY")) {
-      return false;
+      // candidate normalized decay length XY
+      if (candidate.decayLengthXYNormalised() < cuts->get(pTBin, "normDecLXY")) {
+        return false;
+      }
+
+      // candidate impact parameter XY
+      if (std::abs(candidate.impactParameterXY()) > cuts->get(pTBin, "impParXY")) {
+        return false;
+      }
     }
 
-    // candidate impact parameter XY
-    if (std::abs(candidate.impactParameterXY()) > cuts->get(pTBin, "impParXY")) {
-      return false;
+    if constexpr (reconstructionType == aod::hf_cand::VertexerType::KfParticle) {
+      if (applyKfCuts) {
+        // candidate chi2geo of the triplet of prongs
+        if (candidate.kfChi2Geo() > kfCuts->get(pTBin, "kfChi2Geo")) {
+          return false;
+        }
+
+        // candidate's decay point separation from the PV in terms of its error
+        if (candidate.kfDecayLength() / candidate.kfDecayLengthError() < kfCuts->get(pTBin, "kfLdL")) {
+          return false;
+        }
+
+        // candidate's chi2 to the PV
+        if (candidate.kfChi2Topo() > kfCuts->get(pTBin, "kfChi2Topo")) {
+          return false;
+        }
+      }
     }
 
     if (!isSelectedCandidateProngDca(candidate)) {
@@ -242,40 +266,117 @@ struct HfCandidateSelectorLc {
       return false;
     }
 
-    // cut on daughter pT
-    if (trackProton.pt() < cuts->get(pTBin, "pT p") || trackKaon.pt() < cuts->get(pTBin, "pT K") || trackPion.pt() < cuts->get(pTBin, "pT Pi")) {
-      return false;
-    }
+    if (reconstructionType == aod::hf_cand::VertexerType::DCAFitter || (reconstructionType == aod::hf_cand::VertexerType::KfParticle && applyNonKfCuts)) {
 
-    // cut on Lc->pKpi, piKp mass values
-    /// cut on the Kpi pair invariant mass, to study Lc->pK*(->Kpi)
-    float massLc, massKPi;
-    if constexpr (reconstructionType == aod::hf_cand::VertexerType::DCAFitter) {
-      if (trackProton.globalIndex() == candidate.prong0Id()) {
-        massLc = hfHelper.invMassLcToPKPi(candidate);
-        massKPi = hfHelper.invMassKPiPairLcToPKPi(candidate);
-      } else {
-        massLc = hfHelper.invMassLcToPiKP(candidate);
-        massKPi = hfHelper.invMassKPiPairLcToPiKP(candidate);
+      // cut on daughter pT
+      if (trackProton.pt() < cuts->get(pTBin, "pT p") || trackKaon.pt() < cuts->get(pTBin, "pT K") || trackPion.pt() < cuts->get(pTBin, "pT Pi")) {
+        return false;
       }
-    } else if constexpr (reconstructionType == aod::hf_cand::VertexerType::KfParticle) {
-      if (trackProton.globalIndex() == candidate.prong0Id()) {
-        massLc = candidate.kfMassPKPi();
-        massKPi = candidate.kfMassKPi();
-      } else {
-        massLc = candidate.kfMassPiKP();
-        massKPi = candidate.kfMassPiK();
+
+      float massLc, massKPi;
+      if constexpr (reconstructionType == aod::hf_cand::VertexerType::DCAFitter) {
+        if (trackProton.globalIndex() == candidate.prong0Id()) {
+          massLc = hfHelper.invMassLcToPKPi(candidate);
+          massKPi = hfHelper.invMassKPiPairLcToPKPi(candidate);
+        } else {
+          massLc = hfHelper.invMassLcToPiKP(candidate);
+          massKPi = hfHelper.invMassKPiPairLcToPiKP(candidate);
+        }
+      } else if constexpr (reconstructionType == aod::hf_cand::VertexerType::KfParticle) {
+        if (trackProton.globalIndex() == candidate.prong0Id()) {
+          massLc = candidate.kfMassPKPi();
+          massKPi = candidate.kfMassKPi();
+        } else {
+          massLc = candidate.kfMassPiKP();
+          massKPi = candidate.kfMassPiK();
+        }
       }
-    }
 
-    if (std::abs(massLc - o2::constants::physics::MassLambdaCPlus) > cuts->get(pTBin, "m")) {
-      return false;
+      // cut on Lc->pKpi, piKp mass values
+      if (std::abs(massLc - o2::constants::physics::MassLambdaCPlus) > cuts->get(pTBin, "m")) {
+        return false;
+      }
+
+      /// cut on the Kpi pair invariant mass, to study Lc->pK*(->Kpi)
+      const double cutMassKPi = cuts->get(pTBin, "mass (Kpi)");
+      if (cutMassKPi > 0 && std::abs(massKPi - massK0Star892) > cutMassKPi) {
+        return false;
+      }
     }
 
-    /// cut on the Kpi pair invariant mass, to study Lc->pK*(->Kpi)
-    const double cutMassKPi = cuts->get(pTBin, "mass (Kpi)");
-    if (cutMassKPi > 0 && std::abs(massKPi - massK0Star892) > cutMassKPi) {
-      return false;
+    if constexpr (reconstructionType == aod::hf_cand::VertexerType::KfParticle) {
+      if (applyKfCuts) {
+        const float chi2PrimProng0 = candidate.kfChi2PrimProng0();
+        const float chi2PrimProng1 = candidate.kfChi2PrimProng1();
+        const float chi2PrimProng2 = candidate.kfChi2PrimProng2();
+
+        const float chi2GeoProng1Prong2 = candidate.kfChi2GeoProng1Prong2();
+        const float chi2GeoProng0Prong2 = candidate.kfChi2GeoProng0Prong2();
+        const float chi2GeoProng0Prong1 = candidate.kfChi2GeoProng0Prong1();
+
+        const float dcaProng1Prong2 = candidate.kfDcaProng1Prong2();
+        const float dcaProng0Prong2 = candidate.kfDcaProng0Prong2();
+        const float dcaProng0Prong1 = candidate.kfDcaProng0Prong1();
+
+        const double cutChi2PrimPr = kfCuts->get(pTBin, "kfChi2PrimPr");
+        const double cutChi2PrimKa = kfCuts->get(pTBin, "kfChi2PrimKa");
+        const double cutChi2PrimPi = kfCuts->get(pTBin, "kfChi2PrimPi");
+
+        const double cutChi2GeoKaPi = kfCuts->get(pTBin, "kfChi2GeoKaPi");
+        const double cutChi2GeoPrPi = kfCuts->get(pTBin, "kfChi2GeoPrPi");
+        const double cutChi2GeoPrKa = kfCuts->get(pTBin, "kfChi2GeoPrKa");
+
+        const double cutDcaKaPi = kfCuts->get(pTBin, "kfDcaKaPi");
+        const double cutDcaPrPi = kfCuts->get(pTBin, "kfDcaPrPi");
+        const double cutDcaPrKa = kfCuts->get(pTBin, "kfDcaPrKa");
+
+        // kaon's chi2 to the PV
+        if (chi2PrimProng1 < cutChi2PrimKa) {
+          return false;
+        }
+
+        // chi2geo between proton and pion
+        if (chi2GeoProng0Prong2 > cutChi2GeoPrPi) {
+          return false;
+        }
+
+        // dca between proton and pion
+        if (dcaProng0Prong2 > cutDcaPrPi) {
+          return false;
+        }
+
+        const bool isPKPi = trackProton.globalIndex() == candidate.prong0Id();
+
+        // 0-th prong chi2 to the PV
+        if (chi2PrimProng0 < (isPKPi ? cutChi2PrimPr : cutChi2PrimPi)) {
+          return false;
+        }
+
+        // 2-nd prong chi2 to the PV
+        if (chi2PrimProng2 < (isPKPi ? cutChi2PrimPi : cutChi2PrimPr)) {
+          return false;
+        }
+
+        // chi2geo between 1-st and 2-nd prongs
+        if (chi2GeoProng1Prong2 > (isPKPi ? cutChi2GeoKaPi : cutChi2GeoPrKa)) {
+          return false;
+        }
+
+        // chi2geo between 0-th and 1-st prongs
+        if (chi2GeoProng0Prong1 > (isPKPi ? cutChi2GeoPrKa : cutChi2GeoKaPi)) {
+          return false;
+        }
+
+        // dca between 1-st and 2-nd prongs
+        if (dcaProng1Prong2 > (isPKPi ? cutDcaKaPi : cutDcaPrKa)) {
+          return false;
+        }
+
+        // dca between 0-th and 1-st prongs
+        if (dcaProng0Prong1 > (isPKPi ? cutDcaPrKa : cutDcaKaPi)) {
+          return false;
+        }
+      }
     }
 
     return true;
@@ -358,7 +459,7 @@ struct HfCandidateSelectorLc {
       }
 
       // conjugate-independent topological selection
-      if (!selectionTopol(candidate)) {
+      if (!selectionTopol<reconstructionType>(candidate)) {
         hfSelLcCandidate(statusLcToPKPi, statusLcToPiKP);
         if (applyMl) {
           hfMlLcToPKPiCandidate(outputMlLcToPKPi, outputMlLcToPiKP);