Fix candidate selection and cleanup OmegaC task

fcatalan92 · fcatalan92 · commit 6df82e6def15 · 2025-04-15T14:21:54.000+02:00
diff --git a/PWGHF/D2H/Tasks/taskOmegac0ToOmegapi.cxx b/PWGHF/D2H/Tasks/taskOmegac0ToOmegapi.cxx
@@ -12,6 +12,7 @@
 /// \file taskOmegac0ToOmegapi.cxx
 /// \brief OmegaC0 analysis task
 /// \author Yunfan Liu <yunfan.liu@cern.ch>, China University of Geosciences
+/// \author Fabio Catalano <fabio.catalano@cern.ch>, University of Houston
 
 #include <vector>
 
@@ -36,28 +37,33 @@ using namespace o2::framework::expressions;
 struct HfTaskOmegac0ToOmegapi {
   // ML inference
   Configurable<bool> applyMl{"applyMl", false, "Flag to apply ML selections"};
-  Configurable<bool> selectionFlagOmegac0{"selectionFlagOmegac0", false, "Selection Flag for Omegac0 candidates"};
+  Configurable<bool> selectionFlagOmegac0{"selectionFlagOmegac0", true, "Select Omegac0 candidates"};
   Configurable<double> yCandGenMax{"yCandGenMax", 0.5, "max. gen particle rapidity"};
   Configurable<double> yCandRecoMax{"yCandRecoMax", 0.8, "max. cand. rapidity"};
 
   HfHelper hfHelper;
   SliceCache cache;
-  using MyTracksWMc = soa::Join<aod::Tracks, aod::TracksIU, aod::McTrackLabels>;
 
-  using Omegac0Candidates = soa::Join<aod::HfCandToOmegaPi, aod::HfSelToOmegaPi>;
-  using Omegac0CandidatesKF = soa::Join<Omegac0Candidates, aod::HfOmegacKf>;
-  using OmegaC0CandidatesMcKF = soa::Join<Omegac0CandidatesKF, aod::HfToOmegaPiMCRec>;
+  using TracksMc = soa::Join<aod::Tracks, aod::TracksIU, aod::McTrackLabels>;
 
-  using Omegac0CandidatesMl = soa::Join<Omegac0Candidates, aod::HfMlSelOmegacToOmegaPi>;
-  using Omegac0CandidatesMlKF = soa::Join<Omegac0CandidatesMl, aod::HfOmegacKf>;
-  using Omegac0CandidatesMlMcKF = soa::Join<Omegac0CandidatesMlKF, aod::HfToOmegaPiMCRec>;
+  using Omegac0Cands = soa::Filtered<soa::Join<aod::HfCandToOmegaPi, aod::HfSelToOmegaPi>>;
+  using Omegac0CandsKF = soa::Filtered<soa::Join<aod::HfCandToOmegaPi, aod::HfSelToOmegaPi, aod::HfOmegacKf>>;
+  using OmegaC0CandsMcKF = soa::Filtered<soa::Join<aod::HfCandToOmegaPi, aod::HfSelToOmegaPi, aod::HfOmegacKf, aod::HfToOmegaPiMCRec>>;
+
+  using Omegac0CandsMl = soa::Filtered<soa::Join<aod::HfCandToOmegaPi, aod::HfSelToOmegaPi, aod::HfMlSelOmegacToOmegaPi>>;
+  using Omegac0CandsMlKF = soa::Filtered<soa::Join<aod::HfCandToOmegaPi, aod::HfSelToOmegaPi, aod::HfMlSelOmegacToOmegaPi, aod::HfOmegacKf>>;
+  using Omegac0CandsMlMcKF = soa::Filtered<soa::Join<aod::HfCandToOmegaPi, aod::HfSelToOmegaPi, aod::HfMlSelOmegacToOmegaPi, aod::HfOmegacKf, aod::HfToOmegaPiMCRec>>;
+
+  using Omegac0Gen = soa::Filtered<soa::Join<aod::McParticles, aod::HfToOmegaPiMCGen>>;
 
   using Collisions = soa::Join<aod::Collisions, aod::EvSels>;
   using CollisionsWithMcLabels = soa::Join<aod::Collisions, aod::McCollisionLabels, aod::EvSels>;
+
   PresliceUnsorted<CollisionsWithMcLabels> colPerMcCollision = aod::mccollisionlabel::mcCollisionId;
 
-  Partition<Omegac0CandidatesKF> selectedOmegac0CandidatesKF = aod::hf_sel_toomegapi::resultSelections && !selectionFlagOmegac0;
-  Partition<Omegac0CandidatesMlKF> selectedOmegac0CandidatesMlKF = aod::hf_sel_toomegapi::resultSelections && !selectionFlagOmegac0;
+  Filter filterOmegaCToOmegaPiFlag = (aod::hf_track_index::hfflag & static_cast<uint8_t>(BIT(aod::hf_cand_casc_lf::DecayType2Prong::OmegaczeroToOmegaPi))) != static_cast<uint8_t>(0);
+  Filter filterOmegaCMatchedRec = nabs(aod::hf_cand_xic0_omegac0::flagMcMatchRec) == static_cast<int8_t>(BIT(aod::hf_cand_xic0_omegac0::DecayType::OmegaczeroToOmegaPi));
+  Filter filterOmegaCMatchedGen = nabs(aod::hf_cand_xic0_omegac0::flagMcMatchGen) == static_cast<int8_t>(BIT(aod::hf_cand_xic0_omegac0::DecayType::OmegaczeroToOmegaPi));
 
   // ThnSparse for ML outputScores and Vars
   ConfigurableAxis thnConfigAxisPromptScore{"thnConfigAxisPromptScore", {50, 0, 1}, "Prompt score bins"};
@@ -70,18 +76,16 @@ struct HfTaskOmegac0ToOmegapi {
   ConfigurableAxis thnConfigAxisGenPtD{"thnConfigAxisGenPtD", {500, 0, 50}, "Gen Pt D"};
   ConfigurableAxis thnConfigAxisGenPtB{"thnConfigAxisGenPtB", {1000, 0, 100}, "Gen Pt B"};
   ConfigurableAxis thnConfigAxisNumPvContr{"thnConfigAxisNumPvContr", {200, -0.5, 199.5}, "Number of PV contributors"};
-  HistogramRegistry registry{
-    "registry",
-    {}};
+  HistogramRegistry registry{"registry", {}};
 
   void init(InitContext&)
   {
-    std::array<bool, 12> doprocess{doprocessDataWithKFParticle, doprocessMcWithKFParticle, doprocessDataWithKFParticleMl, doprocessMcWithKFParticleMl};
+    std::array<bool, 4> doprocess{doprocessDataWithKFParticle, doprocessMcWithKFParticle, doprocessDataWithKFParticleMl, doprocessMcWithKFParticleMl};
     if ((std::accumulate(doprocess.begin(), doprocess.end(), 0)) != 1) {
       LOGP(fatal, "One and only one process function should be enabled at a time.");
     }
 
-    const AxisSpec thnAxisMass{thnConfigAxisMass, "inv. mass (#Omega #pi) (GeV/#it{c}^{2})"};
+    const AxisSpec thnAxisMass{thnConfigAxisMass, "inv. mass (#Omega#pi) (GeV/#it{c}^{2})"};
     const AxisSpec thnAxisPt{thnConfigAxisPt, "#it{p}_{T} (GeV/#it{c})"};
     const AxisSpec thnAxisPtB{thnConfigAxisPtB, "#it{p}_{T}^{B} (GeV/#it{c})"};
     const AxisSpec thnAxisY{thnConfigAxisY, "y"};
@@ -97,11 +101,7 @@ struct HfTaskOmegac0ToOmegapi {
       registry.get<THnSparse>(HIST("hSparseAcc"))->Sumw2();
     }
 
-    std::vector<AxisSpec> axes = {
-      thnAxisMass,
-      thnAxisPt,
-      thnAxisY,
-    };
+    std::vector<AxisSpec> axes = {thnAxisMass, thnAxisPt, thnAxisY};
     if (doprocessMcWithKFParticle || doprocessMcWithKFParticleMl) {
       axes.push_back(thnAxisPtB);
       axes.push_back(thnAxisOrigin);
@@ -110,9 +110,7 @@ struct HfTaskOmegac0ToOmegapi {
     }
     if (applyMl) {
       const AxisSpec thnAxisPromptScore{thnConfigAxisPromptScore, "BDT score prompt."};
-
       axes.insert(axes.begin(), thnAxisPromptScore);
-
       registry.add("hBdtScoreVsMassVsPtVsPtBVsYVsOriginVsOmegac0Type", "Thn for Omegac0 candidates", HistType::kTHnSparseD, axes);
       registry.get<THnSparse>(HIST("hBdtScoreVsMassVsPtVsPtBVsYVsOriginVsOmegac0Type"))->Sumw2();
     } else {
@@ -125,117 +123,102 @@ struct HfTaskOmegac0ToOmegapi {
   void processData(const CandType& candidates, CollType const&)
   {
     for (const auto& candidate : candidates) {
-      if (!(candidate.hfflag() & 1 << aod::hf_cand_casc_lf::DecayType2Prong::OmegaczeroToOmegaPi)) {
+      if (!(candidate.resultSelections() == true || (candidate.resultSelections() == false && !selectionFlagOmegac0))) {
         continue;
       }
       if (yCandRecoMax >= 0. && std::abs(candidate.kfRapOmegac()) > yCandRecoMax) {
         continue;
       }
-      float massOmegac0;
-      massOmegac0 = candidate.invMassCharmBaryon();
-      auto rapidityCandidate = candidate.kfRapOmegac();
-      auto ptCandidate = candidate.ptCharmBaryon();
+
       if constexpr (applyMl) {
-        registry.fill(HIST("hBdtScoreVsMassVsPtVsPtBVsYVsOriginVsOmegac0Type"), candidate.mlProbOmegac()[0], massOmegac0, ptCandidate, rapidityCandidate);
+        registry.fill(HIST("hBdtScoreVsMassVsPtVsPtBVsYVsOriginVsOmegac0Type"), candidate.mlProbOmegac()[0], candidate.invMassCharmBaryon(), candidate.ptCharmBaryon(), candidate.kfRapOmegac());
       } else {
-        registry.fill(HIST("hMassVsPtVsPtBVsYVsOriginVsOmegac0Type"), massOmegac0, ptCandidate, rapidityCandidate);
+        registry.fill(HIST("hMassVsPtVsPtBVsYVsOriginVsOmegac0Type"), candidate.invMassCharmBaryon(), candidate.ptCharmBaryon(), candidate.kfRapOmegac());
       }
     }
   }
 
-  void processDataWithKFParticle(Omegac0CandidatesKF const&, Collisions const& collisions)
+  void processDataWithKFParticle(Omegac0CandsKF const& candidates, Collisions const& collisions)
   {
-    processData<false>(selectedOmegac0CandidatesKF, collisions);
+    processData<false>(candidates, collisions);
   }
   PROCESS_SWITCH(HfTaskOmegac0ToOmegapi, processDataWithKFParticle, "process HfTaskOmegac0ToOmegapi with KFParticle", false);
-  // TODO: add processKFParticle
 
-  void processDataWithKFParticleMl(Omegac0CandidatesMlKF const&, Collisions const& collisions)
+  void processDataWithKFParticleMl(Omegac0CandsMlKF const& candidates, Collisions const& collisions)
   {
-    processData<true>(selectedOmegac0CandidatesMlKF, collisions);
+    processData<true>(candidates, collisions);
   }
   PROCESS_SWITCH(HfTaskOmegac0ToOmegapi, processDataWithKFParticleMl, "process HfTaskOmegac0ToOmegapi with KFParticle and ML selections", false);
-  // TODO: add processKFParticleMl
 
   template <bool applyMl, typename CandType, typename CollType>
   void processMc(const CandType& candidates,
-                 soa::Join<aod::McParticles, aod::HfToOmegaPiMCGen> const& mcParticles,
-                 MyTracksWMc const&,
+                 Omegac0Gen const& mcParticles,
+                 TracksMc const&,
                  CollType const& collisions,
                  aod::McCollisions const&)
   {
     // MC rec.
     for (const auto& candidate : candidates) {
-      if (!(candidate.hfflag() & 1 << aod::hf_cand_casc_lf::DecayType2Prong::OmegaczeroToOmegaPi)) {
+      if (!(candidate.resultSelections() == true || (candidate.resultSelections() == false && !selectionFlagOmegac0))) {
         continue;
       }
       if (yCandRecoMax >= 0. && std::abs(candidate.kfRapOmegac()) > yCandRecoMax) {
         continue;
       }
-      auto collision = candidate.template collision_as<CollType>();
-      auto numPvContributors = collision.numContrib();
-      float massOmegac0;
-      massOmegac0 = candidate.invMassCharmBaryon();
-      auto ptCandidate = candidate.ptCharmBaryon();
-      auto rapidityCandidate = candidate.kfRapOmegac();
-      if (candidate.resultSelections() && !selectionFlagOmegac0)
-        if (candidate.flagMcMatchRec() == (1 << aod::hf_cand_xic0_omegac0::DecayType::OmegaczeroToOmegaPi)) {
-          if constexpr (applyMl) {
-            registry.fill(HIST("hBdtScoreVsMassVsPtVsPtBVsYVsOriginVsOmegac0Type"), candidate.mlProbOmegac()[0], massOmegac0, ptCandidate, rapidityCandidate, candidate.ptBhadMotherPart(), candidate.originRec(), candidate.flagMcMatchRec(), numPvContributors);
-
-          } else {
-            registry.fill(HIST("hMassVsPtVsPtBVsYVsOriginVsOmegac0Type"), massOmegac0, ptCandidate, rapidityCandidate, candidate.ptBhadMotherPart(), candidate.originRec(), candidate.flagMcMatchRec(), numPvContributors);
-          }
-        }
+
+      auto numPvContributors = candidate.template collision_as<CollType>().numContrib();
+
+      if constexpr (applyMl) {
+        registry.fill(HIST("hBdtScoreVsMassVsPtVsPtBVsYVsOriginVsOmegac0Type"), candidate.mlProbOmegac()[0], candidate.invMassCharmBaryon(), candidate.ptCharmBaryon(), candidate.kfRapOmegac(), candidate.ptBhadMotherPart(), candidate.originRec(), candidate.flagMcMatchRec(), numPvContributors);
+
+      } else {
+        registry.fill(HIST("hMassVsPtVsPtBVsYVsOriginVsOmegac0Type"), candidate.invMassCharmBaryon(), candidate.ptCharmBaryon(), candidate.kfRapOmegac(), candidate.ptBhadMotherPart(), candidate.originRec(), candidate.flagMcMatchRec(), numPvContributors);
+      }
     }
+
     // MC gen.
     for (const auto& particle : mcParticles) {
-      if (std::abs(particle.flagMcMatchGen()) == 1 << aod::hf_cand_xic0_omegac0::DecayType::OmegaczeroToOmegaPi) {
-        if (yCandGenMax >= 0. && std::abs(particle.rapidityCharmBaryonGen()) > yCandGenMax) {
-          continue;
-        }
-        float ptGenB = -1;
-        auto ptGen = particle.pt();
-        auto yGen = particle.rapidityCharmBaryonGen();
-
-        unsigned maxNumContrib = 0;
-        const auto& recoCollsPerMcColl = collisions.sliceBy(colPerMcCollision, particle.mcCollision().globalIndex());
-        for (const auto& recCol : recoCollsPerMcColl) {
-          maxNumContrib = recCol.numContrib() > maxNumContrib ? recCol.numContrib() : maxNumContrib;
-        }
-
-        if (particle.originGen() == RecoDecay::OriginType::Prompt) {
-          registry.fill(HIST("hSparseAcc"), ptGen, ptGenB, yGen, 1, maxNumContrib);
-
-        } else {
-          ptGenB = mcParticles.rawIteratorAt(particle.idxBhadMotherPart()).pt();
-          registry.fill(HIST("hSparseAcc"), ptGen, ptGenB, yGen, 2, maxNumContrib);
-        }
+      if (yCandGenMax >= 0. && std::abs(particle.rapidityCharmBaryonGen()) > yCandGenMax) {
+        continue;
+      }
+
+      auto ptGen = particle.pt();
+      auto yGen = particle.rapidityCharmBaryonGen();
+
+      unsigned maxNumContrib = 0;
+      const auto& recoCollsPerMcColl = collisions.sliceBy(colPerMcCollision, particle.mcCollision().globalIndex());
+      for (const auto& recCol : recoCollsPerMcColl) {
+        maxNumContrib = recCol.numContrib() > maxNumContrib ? recCol.numContrib() : maxNumContrib;
+      }
+
+      if (particle.originGen() == RecoDecay::OriginType::Prompt) {
+        registry.fill(HIST("hSparseAcc"), ptGen, -1., yGen, RecoDecay::OriginType::Prompt, maxNumContrib);
+      } else {
+        float ptGenB = mcParticles.rawIteratorAt(particle.idxBhadMotherPart()).pt();
+        registry.fill(HIST("hSparseAcc"), ptGen, ptGenB, yGen, RecoDecay::OriginType::NonPrompt, maxNumContrib);
       }
     }
   }
 
-  void processMcWithKFParticle(OmegaC0CandidatesMcKF const& omegaC0CandidatesMcKF,
-                               soa::Join<aod::McParticles, aod::HfToOmegaPiMCGen> const& mcParticles,
-                               MyTracksWMc const& tracks,
+  void processMcWithKFParticle(OmegaC0CandsMcKF const& omegaC0CandidatesMcKF,
+                               Omegac0Gen const& mcParticles,
+                               TracksMc const& tracks,
                                CollisionsWithMcLabels const& collisions,
                                aod::McCollisions const& mcCollisions)
   {
     processMc<false>(omegaC0CandidatesMcKF, mcParticles, tracks, collisions, mcCollisions);
   }
   PROCESS_SWITCH(HfTaskOmegac0ToOmegapi, processMcWithKFParticle, "Process MC with KFParticle", false);
-  // TODO: add the processMcWithKFParticle
 
-  void processMcWithKFParticleMl(Omegac0CandidatesMlMcKF const& omegac0CandidatesMlMcKF,
-                                 soa::Join<aod::McParticles, aod::HfToOmegaPiMCGen> const& mcParticles,
-                                 MyTracksWMc const& tracks,
+  void processMcWithKFParticleMl(Omegac0CandsMlMcKF const& omegac0CandidatesMlMcKF,
+                                 Omegac0Gen const& mcParticles,
+                                 TracksMc const& tracks,
                                  CollisionsWithMcLabels const& collisions,
                                  aod::McCollisions const& mcCollisions)
   {
     processMc<true>(omegac0CandidatesMlMcKF, mcParticles, tracks, collisions, mcCollisions);
   }
   PROCESS_SWITCH(HfTaskOmegac0ToOmegapi, processMcWithKFParticleMl, "Process MC with KFParticle and ML selections", false);
-  // TODO: add the processMcWithKFParticleMl
 };
 
 WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)
