[PWGHF] Fix corr. bkg. Sc tree filling. (#15226)

Author: mfaggin

PWGHF/TableProducer/treeCreatorSigmacCorrBkg.cxx

@@ -55,6 +55,8 @@ enum Decays { Sigmac2455Pi = 0,
               LambdacPiPi };
 enum DecaysLambdac { PKPi = 0,
                      PiKP };
+enum Reflections { NotReflected = 0,
+                   Reflected };
 DECLARE_SOA_COLUMN(Y, y, float);
 DECLARE_SOA_COLUMN(Pt, pt, float);
 DECLARE_SOA_COLUMN(Mass, mass, float);
@@ -65,6 +67,7 @@ DECLARE_SOA_COLUMN(Decay, decay, int8_t);
 DECLARE_SOA_COLUMN(DecayLambdac, decayLambdac, int8_t);
 DECLARE_SOA_COLUMN(MlScoreFirstClass, mlScoreFirstClass, float); /// background score Λc
 DECLARE_SOA_COLUMN(MlScoreThirdClass, mlScoreThirdClass, float); /// non-prompt score Λc
+DECLARE_SOA_COLUMN(IsReflected, isReflected, int8_t);
 } // namespace hf_sigmac_bkg
 DECLARE_SOA_TABLE(HfCorrBkgSc, "AOD", "HFCORRBKGSC",
                   hf_sigmac_bkg::Y,
@@ -76,7 +79,8 @@ DECLARE_SOA_TABLE(HfCorrBkgSc, "AOD", "HFCORRBKGSC",
                   hf_sigmac_bkg::Decay,
                   hf_sigmac_bkg::DecayLambdac,
                   hf_sigmac_bkg::MlScoreFirstClass,
-                  hf_sigmac_bkg::MlScoreThirdClass);
+                  hf_sigmac_bkg::MlScoreThirdClass,
+                  hf_sigmac_bkg::IsReflected);
 } // namespace o2::aod
 
 struct HfTreeCreatorSigmacCorrBkg {
@@ -86,6 +90,7 @@ struct HfTreeCreatorSigmacCorrBkg {
   /// Selection of candidates Λc+
   Configurable<int> selectionFlagLc{"selectionFlagLc", 1, "Selection Flag for Lc"};
   Configurable<float> yCandRecoMax{"yCandRecoMax", -1, "Maximum Sc candidate rapidity"};
+  Configurable<bool> keepReflectedSignals{"keepReflectedSignals", false, "Keep also the reflected signals (i.e. generated pKpi reco as piKp, and viceversa)"};
 
   using RecoLcMc = soa::Join<aod::HfCand3Prong, aod::HfCand3ProngMcRec, aod::HfSelLc, aod::HfMlLcToPKPi>;
   using RecoScMc = soa::Join<aod::HfCandSc, aod::HfCandScMcRec>;
@@ -103,33 +108,56 @@ struct HfTreeCreatorSigmacCorrBkg {
     float massLc = -1.f;
     float deltaMass = -1.f;
     const int8_t isCandPKPiPiKP = hf_sigmac_utils::isDecayToPKPiToPiKP(candLcDauSc, candidateSc);
-    std::array<float, 2> outputMl{-1., -1.};
+    std::array<float, 2> outputMlLcPKPi{-1., -1.};
+    std::array<float, 2> outputMlLcPiKP{-1., -1.};
     /// rapidity selection on Σc0,++
     if (yCandRecoMax >= 0. && std::abs(rapidity) > yCandRecoMax) {
       return;
     }
 
     /// BDT scores
+    // reconstructed Λc± → pK-π+
+    if (!candLcDauSc.mlProbLcToPKPi().empty()) {
+      outputMlLcPKPi.at(0) = candLcDauSc.mlProbLcToPKPi()[0]; /// bkg score
+      outputMlLcPKPi.at(1) = candLcDauSc.mlProbLcToPKPi()[2]; /// non-prompt score
+    }
+    // reconstructed Λc± → π+K-p
     if (!candLcDauSc.mlProbLcToPiKP().empty()) {
-      outputMl.at(0) = candLcDauSc.mlProbLcToPiKP()[0]; /// bkg score
-      outputMl.at(1) = candLcDauSc.mlProbLcToPiKP()[2]; /// non-prompt score
+      outputMlLcPiKP.at(0) = candLcDauSc.mlProbLcToPiKP()[0]; /// bkg score
+      outputMlLcPiKP.at(1) = candLcDauSc.mlProbLcToPiKP()[2]; /// non-prompt score
     }
 
-    if ((TESTBIT(isCandPKPiPiKP, o2::aod::hf_cand_sigmac::Decays::PKPi)) && std::abs(candLcDauSc.template prong0_as<aod::TracksWMc>().template mcParticle_as<ParticlesLcSigmac>().pdgCode()) == kProton) {
+    int pdgCodeProng0Abs = std::abs(candLcDauSc.template prong0_as<aod::TracksWMc>().template mcParticle_as<ParticlesLcSigmac>().pdgCode());
+
+    /// candidates with Λc± reconstructed in the pK-π+ decay
+    if ((TESTBIT(isCandPKPiPiKP, o2::aod::hf_cand_sigmac::Decays::PKPi))) {
+
       massSc = HfHelper::invMassScRecoLcToPKPi(candidateSc, candLcDauSc);
       massLc = HfHelper::invMassLcToPKPi(candLcDauSc);
       deltaMass = massSc - massLc;
 
-      /// fill the tree
-      rowCorrBkgSc(rapidity, candidateSc.pt(), massSc, deltaMass, chargeSc, motherPdg, motherDecay, aod::hf_sigmac_bkg::DecaysLambdac::PKPi, outputMl.at(0), outputMl.at(1));
+      if (pdgCodeProng0Abs == kProton) {
+        /// candidates with reconstructed Λc± → pK-π+ decay that are generated Λc± → pK-π+ (not reflected)
+        rowCorrBkgSc(rapidity, candidateSc.pt(), massSc, deltaMass, chargeSc, motherPdg, motherDecay, aod::hf_sigmac_bkg::DecaysLambdac::PKPi, outputMlLcPKPi.at(0), outputMlLcPKPi.at(1), aod::hf_sigmac_bkg::Reflections::NotReflected);
+      } else if (keepReflectedSignals && pdgCodeProng0Abs == kPiPlus) {
+        /// candidates with reconstructed Λc± → pK-π+ decay that are actually generated Λc± → π+K-p (reflected)
+        rowCorrBkgSc(rapidity, candidateSc.pt(), massSc, deltaMass, chargeSc, motherPdg, motherDecay, aod::hf_sigmac_bkg::DecaysLambdac::PKPi, outputMlLcPKPi.at(0), outputMlLcPKPi.at(1), aod::hf_sigmac_bkg::Reflections::Reflected);
+      }
     }
-    if ((TESTBIT(isCandPKPiPiKP, o2::aod::hf_cand_sigmac::Decays::PiKP)) && std::abs(candLcDauSc.template prong0_as<aod::TracksWMc>().template mcParticle_as<ParticlesLcSigmac>().pdgCode()) == kPiPlus) {
+    /// candidates with Λc± reconstructed in the π+K-p decay
+    if ((TESTBIT(isCandPKPiPiKP, o2::aod::hf_cand_sigmac::Decays::PiKP))) {
+
       massSc = HfHelper::invMassScRecoLcToPiKP(candidateSc, candLcDauSc);
       massLc = HfHelper::invMassLcToPiKP(candLcDauSc);
       deltaMass = massSc - massLc;
 
-      /// fill the tree
-      rowCorrBkgSc(rapidity, candidateSc.pt(), massSc, deltaMass, chargeSc, motherPdg, motherDecay, aod::hf_sigmac_bkg::DecaysLambdac::PiKP, outputMl.at(0), outputMl.at(1));
+      if (pdgCodeProng0Abs == kPiPlus) {
+        /// candidates with reconstructed Λc± → π+K-p decay that are generated Λc± → π+K-p (not reflected)
+        rowCorrBkgSc(rapidity, candidateSc.pt(), massSc, deltaMass, chargeSc, motherPdg, motherDecay, aod::hf_sigmac_bkg::DecaysLambdac::PiKP, outputMlLcPiKP.at(0), outputMlLcPiKP.at(1), aod::hf_sigmac_bkg::Reflections::NotReflected);
+      } else if (keepReflectedSignals && pdgCodeProng0Abs == kProton) {
+        /// candidates with reconstructed Λc± → π+K-p decay that are actually generated Λc± → pK-π+ (reflected)
+        rowCorrBkgSc(rapidity, candidateSc.pt(), massSc, deltaMass, chargeSc, motherPdg, motherDecay, aod::hf_sigmac_bkg::DecaysLambdac::PiKP, outputMlLcPiKP.at(0), outputMlLcPiKP.at(1), aod::hf_sigmac_bkg::Reflections::Reflected);
+      }
     }
   }