From e61d9de024a056dec662fc4ced3bbf31ae16ff52 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?V=C3=ADt=20Ku=C4=8Dera?=
 <26327373+vkucera@users.noreply.github.com>
Date: Tue, 29 Sep 2026 16:58:17 +0200
Subject: [PATCH] derivedDataCreatorDsToKKPi: Fix PiKK selection

---
 .../derivedDataCreatorDsToKKPi.cxx            | 19 ++++++++++++++-----
 1 file changed, 14 insertions(+), 5 deletions(-)

diff --git a/PWGHF/TableProducer/derivedDataCreatorDsToKKPi.cxx b/PWGHF/TableProducer/derivedDataCreatorDsToKKPi.cxx
index 7ed944efb1c..b1a2562c47a 100644
--- a/PWGHF/TableProducer/derivedDataCreatorDsToKKPi.cxx
+++ b/PWGHF/TableProducer/derivedDataCreatorDsToKKPi.cxx
@@ -89,7 +89,7 @@ struct HfDerivedDataCreatorDsToKKPi {
 
   SliceCache cache;
   static constexpr double Mass{o2::constants::physics::MassDS};
-  static constexpr int NHypothesesCand{1}; // Number of possible selection hypotheses per candidate.
+  static constexpr int NHypothesesCand{2}; // Number of possible selection hypotheses per candidate.
 
   using CollisionsWCentMult = soa::Join<aod::Collisions, aod::CentFV0As, aod::CentFT0Ms, aod::CentFT0As, aod::CentFT0Cs, aod::PVMultZeqs>;
   using CollisionsWMcCentMult = soa::Join<aod::Collisions, aod::McCollisionLabels, aod::CentFV0As, aod::CentFT0Ms, aod::CentFT0As, aod::CentFT0Cs, aod::PVMultZeqs>;
@@ -101,7 +101,7 @@ struct HfDerivedDataCreatorDsToKKPi {
   using MatchedGenCandidatesMc = soa::Filtered<soa::Join<aod::McParticles, aod::HfCand3ProngMcGen>>;
   using TypeMcCollisions = soa::Join<aod::McCollisions, aod::McCentFT0Ms>;
 
-  Filter filterSelectCandidates = (aod::hf_sel_candidate_ds::isSelDsToKKPi & static_cast<int>(BIT(aod::SelectionStep::RecoMl - 1))) != 0; // select candidates which passed all cuts at least up to RecoMl - 1
+  Filter filterSelectCandidates = (aod::hf_sel_candidate_ds::isSelDsToKKPi & static_cast<int>(BIT(aod::SelectionStep::RecoMl - 1))) != 0 || (aod::hf_sel_candidate_ds::isSelDsToPiKK & static_cast<int>(BIT(aod::SelectionStep::RecoMl - 1))) != 0; // select candidates which passed all cuts at least up to RecoMl - 1
   Filter filterMcGenMatching = (nabs(aod::hf_cand_mc_flag::flagMcMatchGen) == static_cast<int8_t>(hf_decay::hf_cand_3prong::DecayChannelMain::DsToPiKK)) || (acceptCorrelatedBkgs && aod::hf_cand_mc_flag::flagMcMatchGen != 0);
 
   Preslice<SelectedCandidates> candidatesPerCollision = aod::hf_cand::collisionId;
@@ -266,7 +266,7 @@ struct HfDerivedDataCreatorDsToKKPi {
       int8_t flagMcRec = 0, origin = 0, swapping = 0, flagDecayChanRec = 0;
       for (const auto& candidate : candidatesThisColl) {
         if constexpr (IsMl) {
-          if (!TESTBIT(candidate.isSelDsToKKPi(), aod::SelectionStep::RecoMl)) {
+          if (!TESTBIT(candidate.isSelDsToKKPi(), aod::SelectionStep::RecoMl) && !TESTBIT(candidate.isSelDsToPiKK(), aod::SelectionStep::RecoMl)) {
             continue;
           }
         }
@@ -295,11 +295,20 @@ struct HfDerivedDataCreatorDsToKKPi {
         double const ct = HfHelper::ctDs(candidate);
         double const y = HfHelper::yDs(candidate);
         float const massDsToKKPi = HfHelper::invMassDsToKKPi(candidate);
-        std::vector<float> mlScoresDsToKKPi;
+        float const massDsToPiKK = HfHelper::invMassDsToPiKK(candidate);
+        std::vector<float> mlScoresDsToKKPi, mlScoresDsToPiKK;
+        auto selectionStep = aod::SelectionStep::RecoMl - 1;
         if constexpr (IsMl) {
+          selectionStep = aod::SelectionStep::RecoMl;
           std::copy(candidate.mlProbDsToKKPi().begin(), candidate.mlProbDsToKKPi().end(), std::back_inserter(mlScoresDsToKKPi));
+          std::copy(candidate.mlProbDsToPiKK().begin(), candidate.mlProbDsToPiKK().end(), std::back_inserter(mlScoresDsToPiKK));
+        }
+        if (TESTBIT(candidate.isSelDsToKKPi(), selectionStep)) {
+          fillTablesCandidate(candidate, 0, massDsToKKPi, ct, y, flagMcRec, origin, swapping, flagDecayChanRec, mlScoresDsToKKPi);
+        }
+        if (TESTBIT(candidate.isSelDsToPiKK(), selectionStep)) {
+          fillTablesCandidate(candidate, 1, massDsToPiKK, ct, y, flagMcRec, origin, swapping, flagDecayChanRec, mlScoresDsToPiKK);
         }
-        fillTablesCandidate(candidate, 0, massDsToKKPi, ct, y, flagMcRec, origin, swapping, flagDecayChanRec, mlScoresDsToKKPi);
       }
     }
   }