[PWGLF] store wd mother pdg info in mask (#9285)

Co-authored-by: ALICE Action Bot <[email protected]>

PWGLF/TableProducer/Nuspex/ebyeMaker.cxx

@@ -159,6 +159,7 @@ struct CandidateTrack {
   float genpt = -999.f;
   float geneta = -999.f;
   int pdgcode = -999;
+  int pdgcodemoth = -999;
   bool isreco = 0;
   int64_t mcIndex = -999;
   int64_t globalIndex = -999;
@@ -179,6 +180,12 @@ enum selBits {
   kTPCPIDMid = BIT(11)
 };
 
+enum PartTypes {
+  kLa = BIT(20),
+  kSig = BIT(21),
+  kPhysPrim = BIT(22)
+};
+
 struct tagRun2V0MCalibration {
   bool mCalibrationStored = false;
   TH1* mhVtxAmpCorrV0A = nullptr;
@@ -313,6 +320,37 @@ struct ebyeMaker {
   Preslice<aod::V0s> perCollisionV0 = o2::aod::v0::collisionId;
   Preslice<aod::McParticles> perCollisionMcParts = o2::aod::mcparticle::mcCollisionId;
 
+  template <class P>
+  int getPartTypeMother(P const& mcPart)
+  {
+    for (auto& mother : mcPart.template mothers_as<aod::McParticles>()) {
+      if (!mother.isPhysicalPrimary())
+        return -1;
+      int pdgCode = mother.pdgCode();
+      switch (std::abs(pdgCode)) {
+        case 3122: {
+          int foundPi = 0;
+          for (auto& mcDaught : mother.template daughters_as<aod::McParticles>()) {
+            if (std::abs(mcDaught.pdgCode()) == 211) {
+              foundPi = mcDaught.pdgCode();
+              break;
+            }
+          }
+          if (foundPi * mcPart.pdgCode() < -0.5)
+            return PartTypes::kLa;
+          return -1;
+        }
+        // case 3222:
+        //   return PartTypes::kSig;
+        // case 3112:
+        //   return PartTypes::kSig;
+        default:
+          return -1;
+      }
+    }
+    return -1;
+  }
+
   template <class T>
   bool selectV0Daughter(T const& track)
   {
@@ -867,7 +905,7 @@ struct ebyeMaker {
   }
 
   template <class C, class T>
-  void fillMcEvent(C const& collision, T const& tracks, aod::V0s const& V0s, float const& centrality, aod::McParticles const&, aod::McTrackLabels const& mcLabels)
+  void fillMcEvent(C const& collision, T const& tracks, aod::V0s const& V0s, float const& centrality, aod::McParticles const& particlesMC, aod::McTrackLabels const& mcLabels)
   {
     fillRecoEvent<C, T>(collision, tracks, V0s, centrality);
 
@@ -876,14 +914,23 @@ struct ebyeMaker {
         candidateTrack.isreco = true;
 
         auto mcLab = mcLabels.rawIteratorAt(candidateTrack.globalIndex);
+
+        if (mcLab.mcParticleId() < -1 || mcLab.mcParticleId() >= particlesMC.size()) {
+          continue;
+        }
         if (mcLab.has_mcParticle()) {
           auto mcTrack = mcLab.template mcParticle_as<aod::McParticles>();
           if (std::abs(mcTrack.pdgCode()) != partPdg[iP])
             continue;
-          if (((mcTrack.flags() & 0x8) && (doprocessMcRun2 || doprocessMiniMcRun2)) || (mcTrack.flags() & 0x2) || (mcTrack.flags() & 0x1))
+          if (((mcTrack.flags() & 0x8) && (doprocessMcRun2 || doprocessMiniMcRun2)) || (mcTrack.flags() & 0x2) || ((mcTrack.flags() & 0x1) && !doprocessMiniMcRun2))
             continue;
-          if (!mcTrack.isPhysicalPrimary())
+
+          if (!mcTrack.isPhysicalPrimary() && !doprocessMiniMcRun2)
             continue;
+          if (mcTrack.isPhysicalPrimary())
+            candidateTrack.pdgcodemoth = PartTypes::kPhysPrim;
+          else if (mcTrack.has_mothers() && iP == 0)
+            candidateTrack.pdgcodemoth = getPartTypeMother(mcTrack);
 
           auto genPt = std::hypot(mcTrack.px(), mcTrack.py());
           candidateTrack.pdgcode = mcTrack.pdgCode();
@@ -936,7 +983,7 @@ struct ebyeMaker {
       if (std::abs(genEta) > etaMax) {
         continue;
       }
-      if (((mcPart.flags() & 0x8) && (doprocessMcRun2 || doprocessMiniMcRun2)) || (mcPart.flags() & 0x2) || (mcPart.flags() & 0x1))
+      if (((mcPart.flags() & 0x8) && (doprocessMcRun2 || doprocessMiniMcRun2)) || (mcPart.flags() & 0x2) || ((mcPart.flags() & 0x1) && !doprocessMiniMcRun2))
         continue;
       auto pdgCode = mcPart.pdgCode();
       if (std::abs(pdgCode) == 3122) {
@@ -969,13 +1016,18 @@ struct ebyeMaker {
         if (std::abs(pdgCode) == partPdg[0]) {
           iP = 0;
         }
-        if (!mcPart.isPhysicalPrimary() && !mcPart.has_mothers())
+        if ((!mcPart.isPhysicalPrimary() && !doprocessMiniMcRun2))
           continue;
         auto genPt = std::hypot(mcPart.px(), mcPart.py());
         CandidateTrack candTrack;
         candTrack.genpt = genPt;
         candTrack.geneta = mcPart.eta();
         candTrack.pdgcode = pdgCode;
+        if (mcPart.isPhysicalPrimary())
+          candTrack.pdgcodemoth = PartTypes::kPhysPrim;
+        else if (mcPart.has_mothers() && iP == 0)
+          candTrack.pdgcodemoth = getPartTypeMother(mcPart);
+
         auto it = find_if(candidateTracks[iP].begin(), candidateTracks[iP].end(), [&](CandidateTrack trk) { return trk.mcIndex == mcPart.globalIndex(); });
         if (it != candidateTracks[iP].end()) {
           continue;
@@ -1385,7 +1437,13 @@ struct ebyeMaker {
           selMask = getTrackSelMask(candidateTrack);
           // if (candidateTrack.outerPID < -4)
           //   continue;
+          if (candidateTrack.pdgcodemoth > 0)
+            selMask |= candidateTrack.pdgcodemoth;
+        } else if (candidateTrack.pdgcodemoth > 0) {
+          selMask = candidateTrack.pdgcodemoth;
         }
+        if (selMask < 0)
+          continue;
         mcMiniTrkTable(
           miniCollTable.lastIndex(),
           candidateTrack.pt,