[PWGCF/FemtoUniverse] Properly filter by PDG and charge in track-track-extended task (AliceO2Group#10401)

Author: davkk

PWGCF/FemtoUniverse/Tasks/femtoUniversePairTaskTrackTrackExtended.cxx

@@ -16,6 +16,7 @@
 /// \author Anton Riedel, TU München, [email protected]
 /// \author Zuzanna Chochulska, WUT Warsaw & CTU Prague, [email protected]
 
+#include <Framework/O2DatabasePDGPlugin.h>
 #include <string>
 #include <vector>
 
@@ -53,6 +54,8 @@ static const float cutsTable[NPart][NCuts]{
 } // namespace
 
 struct FemtoUniversePairTaskTrackTrackExtended {
+  Service<o2::framework::O2DatabasePDG> pdgMC;
+
   /// Particle selection part
 
   /// Table for both particles
@@ -97,7 +100,6 @@ struct FemtoUniversePairTaskTrackTrackExtended {
 
   Partition<soa::Join<FilteredFemtoFullParticles, aod::FDMCLabels>> partsOneMCTruth =
     aod::femtouniverseparticle::partType == static_cast<uint8_t>(aod::femtouniverseparticle::ParticleType::kMCTruthTrack) &&
-    aod::femtouniverseparticle::pidCut == static_cast<uint32_t>(trackonefilter.confPDGCodePartOne) &&
     aod::femtouniverseparticle::pt < trackonefilter.confPtHighPart1 &&
     aod::femtouniverseparticle::pt > trackonefilter.confPtLowPart1;
 
@@ -122,7 +124,6 @@ struct FemtoUniversePairTaskTrackTrackExtended {
 
   Partition<soa::Join<FilteredFemtoFullParticles, aod::FDMCLabels>> partsTwoMCTruth =
     aod::femtouniverseparticle::partType == static_cast<uint8_t>(aod::femtouniverseparticle::ParticleType::kMCTruthTrack) &&
-    aod::femtouniverseparticle::pidCut == static_cast<uint32_t>(tracktwofilter.confPDGCodePartTwo) &&
     aod::femtouniverseparticle::pt < tracktwofilter.confPtHighPart2 &&
     aod::femtouniverseparticle::pt > tracktwofilter.confPtLowPart2;
 
@@ -313,6 +314,12 @@ struct FemtoUniversePairTaskTrackTrackExtended {
     return false;
   }
 
+  /// @returns 1 if positive, -1 if negative, 0 if zero
+  auto sign(auto number) -> int8_t
+  {
+    return (number > 0) - (number < 0);
+  }
+
   void init(InitContext&)
   {
     if (twotracksconfigs.confIsMC) {
@@ -353,6 +360,38 @@ struct FemtoUniversePairTaskTrackTrackExtended {
     eventHisto.fillQA(col);
   }
 
+  template <uint8_t N>
+    requires isOneOrTwo<N>
+  auto doMCTruth(FemtoUniverseParticleHisto<aod::femtouniverseparticle::ParticleType::kMCTruthTrack, N> hist, auto parts) -> void
+  {
+    auto expectedPDG = 0;
+    auto expectedCharge = 0.0l;
+
+    if constexpr (N == ParticleNo::ONE) {
+      expectedPDG = trackonefilter.confPDGCodePartOne;
+      expectedCharge = trackonefilter.confChargePart1;
+    } else if constexpr (N == ParticleNo::TWO) {
+      expectedPDG = tracktwofilter.confPDGCodePartTwo;
+      expectedCharge = tracktwofilter.confChargePart2;
+    }
+
+    for (const auto& particle : parts) {
+      auto pdgCode = static_cast<int>(particle.pidCut());
+      if (pdgCode != expectedPDG) {
+        continue;
+      }
+
+      const auto& pdgParticle = pdgMC->GetParticle(pdgCode);
+      if (!pdgParticle) {
+        continue;
+      }
+
+      if (sign(pdgParticle->Charge()) == sign(expectedCharge)) {
+        hist.template fillQA<false, false>(particle);
+      }
+    }
+  }
+
   /// This function processes the same event and takes care of all the histogramming
   /// \todo the trivial loops over the tracks should be factored out since they will be common to all combinations of T-T, T-V0, V0-V0, ...
   /// @tparam PartitionType
@@ -563,15 +602,11 @@ struct FemtoUniversePairTaskTrackTrackExtended {
     fillCollision(col);
 
     auto groupMCTruth1 = partsOneMCTruth->sliceByCached(aod::femtouniverseparticle::fdCollisionId, col.globalIndex(), cache);
-    for (const auto& particle : groupMCTruth1) {
-      hMCTruth1.fillQA<false, false>(particle);
-    }
+    doMCTruth<1>(hMCTruth1, groupMCTruth1);
 
     if (!confIsSame) {
       auto groupMCTruth2 = partsTwoMCTruth->sliceByCached(aod::femtouniverseparticle::fdCollisionId, col.globalIndex(), cache);
-      for (const auto& particle : groupMCTruth1) {
-        hMCTruth2.fillQA<false, false>(particle);
-      }
+      doMCTruth<2>(hMCTruth2, groupMCTruth2);
     }
 
     auto groupMCReco1 = partsOneMCReco->sliceByCached(aod::femtouniverseparticle::fdCollisionId, col.globalIndex(), cache);