PWGLF: add gen cascades table to non-prompt-cascades task (#8115)

AliceO2Group · Oct 23, 2024 · e9e32ec · e9e32ec
1 parent 256cbcb
commit e9e32ec
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 15 deletions.
diff --git a/PWGLF/DataModel/LFNonPromptCascadeTables.h b/PWGLF/DataModel/LFNonPromptCascadeTables.h
@@ -226,6 +226,18 @@ DECLARE_SOA_TABLE(NPCascTableMC, "AOD", "NPCASCTABLEMC",
                   NPCascadeTable::DCAzMC,
                   NPCascadeTable::MCcollisionMatch)
 
+DECLARE_SOA_TABLE(NPCascTableGen, "AOD", "NPCASCTABLEGen",
+                  NPCascadeTable::gPt,
+                  NPCascadeTable::gEta,
+                  NPCascadeTable::gPhi,
+                  NPCascadeTable::PDGcode,
+                  NPCascadeTable::PdgCodeMom,
+                  NPCascadeTable::DCAxMC,
+                  NPCascadeTable::DCAyMC,
+                  NPCascadeTable::DCAzMC,
+                  NPCascadeTable::IsFromBeauty,
+                  NPCascadeTable::IsFromCharm)
+
 } // namespace o2::aod
 
 #endif // PWGLF_DATAMODEL_LFNONPROMPTCASCADETABLES_H_
diff --git a/PWGLF/Tasks/Strangeness/nonPromptCascade.cxx b/PWGLF/Tasks/Strangeness/nonPromptCascade.cxx
@@ -145,6 +145,7 @@ struct NonPromptCascadeTask {
 
   Produces<o2::aod::NPCascTable> NPCTable;
   Produces<o2::aod::NPCascTableMC> NPCTableMC;
+  Produces<o2::aod::NPCascTableGen> NPCTableGen;
 
   using TracksExtData = soa::Join<aod::TracksIU, aod::TracksCovIU, aod::TracksExtra, aod::pidTPCFullKa, aod::pidTPCFullPi, aod::pidTPCFullPr, aod::pidTOFFullKa, aod::pidTOFFullPi, aod::pidTOFFullPr>;
   using TracksExtMC = soa::Join<aod::TracksIU, aod::TracksCovIU, aod::TracksExtra, aod::McTrackLabels, aod::pidTPCFullKa, aod::pidTPCFullPi, aod::pidTPCFullPr, aod::pidTOFFullKa, aod::pidTOFFullPi, aod::pidTOFFullPr>;
@@ -371,6 +372,26 @@ struct NonPromptCascadeTask {
 
     candidates.clear();
     std::vector<int> mcParticleId;
+
+    auto isFromHF = [&](auto particle) -> std::tuple<bool, bool> {
+      bool fromBeauty = false;
+      bool fromCharm = false;
+      if (particle.has_mothers()) {
+        auto mom = particle.template mothers_as<aod::McParticles>()[0];
+        int pdgCodeMom = mom.pdgCode();
+        fromBeauty = std::abs(pdgCodeMom) / 5000 == 1 || std::abs(pdgCodeMom) / 500 == 1 || std::abs(pdgCodeMom) == 5;
+        fromCharm = std::abs(pdgCodeMom) / 4000 == 1 || std::abs(pdgCodeMom) / 400 == 1 || std::abs(pdgCodeMom) == 4;
+        while (mom.has_mothers()) {
+          const auto grandma = mom.template mothers_as<aod::McParticles>()[0];
+          int pdgCodeGrandma = std::abs(grandma.pdgCode());
+          fromBeauty = fromBeauty || (pdgCodeGrandma / 5000 == 1 || pdgCodeGrandma / 500 == 1 || pdgCodeGrandma == 5);
+          fromCharm = fromCharm || (pdgCodeGrandma / 4000 == 1 || pdgCodeGrandma / 400 == 1 || pdgCodeGrandma == 4);
+          mom = grandma;
+        }
+      }
+      return {fromCharm, fromBeauty};
+    };
+
     for (const auto& trackedCascade : trackedCascades) {
       auto collision = trackedCascade.collision_as<CollisionCandidatesRun3MC>();
 
@@ -566,24 +587,13 @@ struct NonPromptCascadeTask {
       bool isGoodMatch = ((motherParticleID == ITStrack.mcParticleId())) ? true : false;
 
       int pdgCodeMom = 0;
-      bool fromBeauty = false, fromCharm = false;
+      std::tuple<bool, bool> fromHF{false, false};
       if (isGoodCascade && isGoodMatch) {
-        if (track.mcParticle().has_mothers()) {
-          auto mom = track.mcParticle().mothers_as<aod::McParticles>()[0];
-          pdgCodeMom = mom.pdgCode();
-          fromBeauty = std::abs(pdgCodeMom) / 5000 == 1 || std::abs(pdgCodeMom) / 500 == 1 || std::abs(pdgCodeMom) == 5;
-          fromCharm = std::abs(pdgCodeMom) / 4000 == 1 || std::abs(pdgCodeMom) / 400 == 1 || std::abs(pdgCodeMom) == 4;
-          while (mom.has_mothers()) {
-            const auto grandma = mom.mothers_as<aod::McParticles>()[0];
-            int pdgCodeGrandma = std::abs(grandma.pdgCode());
-            fromBeauty = fromBeauty || (pdgCodeGrandma / 5000 == 1 || pdgCodeGrandma / 500 == 1 || pdgCodeGrandma == 5);
-            fromCharm = fromCharm || (pdgCodeGrandma / 4000 == 1 || pdgCodeGrandma / 400 == 1 || pdgCodeGrandma == 4);
-            mom = grandma;
-          }
-        }
+        fromHF = isFromHF(track.mcParticle());
+        pdgCodeMom = track.mcParticle().has_mothers() ? track.mcParticle().mothers_as<aod::McParticles>()[0].pdgCode() : 0;
       }
 
-      candidates.emplace_back(NPCascCandidate{track.globalIndex(), ITStrack.globalIndex(), trackedCascade.collisionId(), trackedCascade.matchingChi2(), trackedCascade.itsClsSize(), isGoodMatch, isGoodCascade, pdgCodeMom, fromBeauty, fromCharm,
+      candidates.emplace_back(NPCascCandidate{track.globalIndex(), ITStrack.globalIndex(), trackedCascade.collisionId(), trackedCascade.matchingChi2(), trackedCascade.itsClsSize(), isGoodMatch, isGoodCascade, pdgCodeMom, std::get<0>(fromHF), std::get<1>(fromHF),
                                               primaryVertex.getX(), primaryVertex.getY(), primaryVertex.getZ(),
                                               track.pt(), track.eta(), track.phi(),
                                               protonTrack.pt(), protonTrack.eta(), pionTrack.pt(), pionTrack.eta(), bachelor.pt(), bachelor.eta(),
@@ -626,6 +636,18 @@ struct NonPromptCascadeTask {
                  c.protonTOFNSigma, c.pionTOFNSigma, c.bachKaonTOFNSigma, c.bachPionTOFNSigma,
                  particle.pt(), particle.eta(), particle.phi(), particle.pdgCode(), mcCollision.posX() - particle.vx(), mcCollision.posY() - particle.vy(), mcCollision.posZ() - particle.vz(), mcCollision.globalIndex() == label.mcCollisionId());
     }
+
+    for (auto& p : mcParticles) {
+      auto absCode = std::abs(p.pdgCode());
+      if (absCode != 3312 && absCode != 3334) {
+        continue;
+      }
+      auto fromHF = isFromHF(p);
+      int pdgCodeMom = p.has_mothers() ? p.mothers_as<aod::McParticles>()[0].pdgCode() : 0;
+      auto mcCollision = p.mcCollision_as<aod::McCollisions>();
+
+      NPCTableGen(p.pt(), p.eta(), p.phi(), p.pdgCode(), pdgCodeMom, mcCollision.posX() - p.vx(), mcCollision.posY() - p.vy(), mcCollision.posZ() - p.vz(), std::get<0>(fromHF), std::get<1>(fromHF));
+    }
   }
   PROCESS_SWITCH(NonPromptCascadeTask, processTrackedCascadesMC, "process cascades from strangeness tracking: MC analysis", true);