resolve conflict kfPV: apply mass and topo constraints

Author: lubynets

PWGHF/TableProducer/candidateCreator3Prong.cxx

@@ -85,6 +85,9 @@ struct HfCandidateCreator3Prong {
   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 @@ struct HfCandidateCreator3Prong {
       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,16 @@ struct HfCandidateCreator3Prong {
       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
+        kfCandPKPi.SetProductionVertex(KFPV);
+        kfCandPiKP.SetProductionVertex(KFPV);
+        kfCandPiKPi.SetProductionVertex(KFPV);
+        kfCandKKPi.SetProductionVertex(KFPV);
+        kfCandPiKK.SetProductionVertex(KFPV);
+      }
+
       KFParticle kfPairKPi;
       const KFParticle* kfDaughtersKPi[3] = {&kfSecondKaon, &kfThirdPion};
       kfPairKPi.SetConstructMethod(2);
@@ -505,8 +522,15 @@ struct HfCandidateCreator3Prong {
       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);