Add two MC particle variables

Author: Mingze129

PWGHF/DataModel/DerivedTables.h

@@ -162,6 +162,7 @@ DECLARE_SOA_INDEX_COLUMN(McParticle, mcParticle);                   //! MC parti
 DECLARE_SOA_COLUMN(FlagMcMatchGen, flagMcMatchGen, int8_t);         //! flag for generator level matching
 DECLARE_SOA_COLUMN(OriginMcGen, originMcGen, int8_t);               //! particle origin, generator level
 DECLARE_SOA_COLUMN(FlagMcDecayChanGen, flagMcDecayChanGen, int8_t); //! resonant decay channel flag, generator level
+DECLARE_SOA_COLUMN(PdgMother, pdgMother, int32_t);                  //! PDG code of the mother particle
 } // namespace hf_mc_particle
 
 // Declares the base table with candidates (Bases).
@@ -492,6 +493,7 @@ DECLARE_SOA_COLUMN(MlScoreBkgCharm, mlScoreBkgCharm, float);                  //
 DECLARE_SOA_COLUMN(MlScorePromptCharm, mlScorePromptCharm, float);            //! ML score for prompt class
 DECLARE_SOA_COLUMN(MlScoreNonPromptCharm, mlScoreNonPromptCharm, float);      //! ML score for non-prompt class
 DECLARE_SOA_COLUMN(MlScoresCharm, mlScoresCharm, std::vector<float>);         //! vector of ML scores
+DECLARE_SOA_COLUMN(FlagMcMatchGenCharm, flagMcMatchGenCharm, int8_t);         //! flag for generator level matching
 } // namespace hf_cand_mc_charm
 
 // ----------------
@@ -935,6 +937,11 @@ DECLARE_SOA_TABLE_STAGED(HfDstarMcs, "HFDSTMC", //! Table with MC candidate info
                          hf_cand::NTracksDecayed,
                          o2::soa::Marker<MarkerDstar>);
 
+DECLARE_SOA_TABLE_STAGED(HfDstarGenMcs, "HFDSTGENMC", //! Table with MC candidate info for generated D*+
+                         hf_cand_mc_charm::FlagMcMatchGenCharm,
+                         hf_mc_particle::PdgMother,
+                         o2::soa::Marker<MarkerDstar>);
+
 // ----------------
 // Ξc± → (Ξ∓ → (Λ → p π∓) π∓) π± π±
 // ----------------

PWGHF/TableProducer/derivedDataCreatorDstarToD0Pi.cxx

@@ -70,6 +70,7 @@ struct HfDerivedDataCreatorDstarToD0Pi {
   Produces<o2::aod::HfDstarMls> rowCandidateMl;
   Produces<o2::aod::HfDstarIds> rowCandidateId;
   Produces<o2::aod::HfDstarMcs> rowCandidateMc;
+  Produces<o2::aod::HfDstarGenMcs> rowCandidateGenMc;
 
   // Switches for filling tables
   HfConfigurableDerivedData confDerData;
@@ -79,6 +80,7 @@ struct HfDerivedDataCreatorDstarToD0Pi {
   Configurable<bool> fillCandidateMl{"fillCandidateMl", true, "Fill candidate selection ML scores"};
   Configurable<bool> fillCandidateId{"fillCandidateId", true, "Fill original indices from the candidate table"};
   Configurable<bool> fillCandidateMc{"fillCandidateMc", true, "Fill candidate MC info"};
+  Configurable<bool> fillCandidateGenMc{"fillCandidateGenMc", true, "Fill candidate generator level MC info"};
   // Parameters for production of training samples
   Configurable<float> downSampleBkgFactor{"downSampleBkgFactor", 1., "Fraction of background candidates to keep for ML trainings"};
   Configurable<float> ptMaxForDownSample{"ptMaxForDownSample", 10., "Maximum pt for the application of the downsampling factor"};
@@ -204,6 +206,42 @@ struct HfDerivedDataCreatorDstarToD0Pi {
     }
   }
 
+  template <typename CollisionType, typename ParticleType, typename MatchMap>
+  void processGenMc(CollisionType const& mcCollisions,
+                          o2::framework::Preslice<ParticleType> const& mcParticlesPerMcCollision,
+                          ParticleType const& mcParticles,
+                          MatchMap const& matchMap)
+  {
+    // Fill MC collision properties
+    for (const auto& mcCollision : mcCollisions) {
+      auto thisMcCollId = mcCollision.globalIndex();
+      auto particlesThisMcColl = mcParticles.sliceBy(mcParticlesPerMcCollision, thisMcCollId);
+      auto sizeTablePart = particlesThisMcColl.size();
+      LOGF(debug, "MC collision %d has %d MC particles", thisMcCollId, sizeTablePart);
+      // Skip MC collisions without HF particles (and without HF candidates in matched reconstructed collisions if saving indices of reconstructed collisions matched to MC collisions)
+      bool isEmpty = (matchMap.find(thisMcCollId) == matchMap.end()) || matchMap.find(thisMcCollId)->second.empty();
+      if (sizeTablePart == 0 && isEmpty) {
+        LOGF(debug, "Skipping MC collision %d", thisMcCollId);
+        continue;
+      }
+      reserveTable(rowCandidateGenMc, fillCandidateGenMc, sizeTablePart);
+      if (fillCandidateGenMc) {
+        for (const auto& particle : particlesThisMcColl) {
+          auto origin = particle.originMcGen();
+          int pdgMother = 0;
+          if (origin == RecoDecay::OriginType::NonPrompt) {
+            auto bHadMother = mcParticles.rawIteratorAt(particle.idxBhadMotherPart() - particlesThisMcColl.offset());
+            pdgMother = std::abs(bHadMother.pdgCode());
+          }
+          rowCandidateGenMc(
+            particle.flagMcMatchGenD0(),
+            pdgMother);
+        }
+      }
+      
+    }
+  }
+
   template <bool isMl, bool isMc, bool onlyBkg, bool onlySig, typename CollType, typename CandType>
   void processCandidates(CollType const& collisions,
                          Partition<CandType>& candidates,
@@ -322,6 +360,7 @@ struct HfDerivedDataCreatorDstarToD0Pi {
     rowsCommon.preProcessMcCollisions(mcCollisions, mcParticlesPerMcCollision, mcParticles);
     processCandidates<false, true, false, true>(collisions, candidatesMcSig, tracks, bcs);
     rowsCommon.processMcParticles(mcCollisions, mcParticlesPerMcCollision, mcParticles, Mass);
+    processGenMc(mcCollisions, mcParticlesPerMcCollision, mcParticles, rowsCommon.matchedCollisions);
   }
   PROCESS_SWITCH(HfDerivedDataCreatorDstarToD0Pi, processMcSig, "Process MC only for signals", false);
 
@@ -335,6 +374,7 @@ struct HfDerivedDataCreatorDstarToD0Pi {
     rowsCommon.preProcessMcCollisions(mcCollisions, mcParticlesPerMcCollision, mcParticles);
     processCandidates<false, true, true, false>(collisions, candidatesMcBkg, tracks, bcs);
     rowsCommon.processMcParticles(mcCollisions, mcParticlesPerMcCollision, mcParticles, Mass);
+    processGenMc(mcCollisions, mcParticlesPerMcCollision, mcParticles, rowsCommon.matchedCollisions);
   }
   PROCESS_SWITCH(HfDerivedDataCreatorDstarToD0Pi, processMcBkg, "Process MC only for background", false);
 
@@ -348,6 +388,7 @@ struct HfDerivedDataCreatorDstarToD0Pi {
     rowsCommon.preProcessMcCollisions(mcCollisions, mcParticlesPerMcCollision, mcParticles);
     processCandidates<false, true, false, false>(collisions, candidatesMcAll, tracks, bcs);
     rowsCommon.processMcParticles(mcCollisions, mcParticlesPerMcCollision, mcParticles, Mass);
+    processGenMc(mcCollisions, mcParticlesPerMcCollision, mcParticles, rowsCommon.matchedCollisions);
   }
   PROCESS_SWITCH(HfDerivedDataCreatorDstarToD0Pi, processMcAll, "Process MC", false);
 
@@ -372,6 +413,7 @@ struct HfDerivedDataCreatorDstarToD0Pi {
     rowsCommon.preProcessMcCollisions(mcCollisions, mcParticlesPerMcCollision, mcParticles);
     processCandidates<true, true, false, true>(collisions, candidatesMcMlSig, tracks, bcs);
     rowsCommon.processMcParticles(mcCollisions, mcParticlesPerMcCollision, mcParticles, Mass);
+    processGenMc(mcCollisions, mcParticlesPerMcCollision, mcParticles, rowsCommon.matchedCollisions);
   }
   PROCESS_SWITCH(HfDerivedDataCreatorDstarToD0Pi, processMcMlSig, "Process MC with ML only for signals", false);
 
@@ -385,6 +427,7 @@ struct HfDerivedDataCreatorDstarToD0Pi {
     rowsCommon.preProcessMcCollisions(mcCollisions, mcParticlesPerMcCollision, mcParticles);
     processCandidates<true, true, true, false>(collisions, candidatesMcMlBkg, tracks, bcs);
     rowsCommon.processMcParticles(mcCollisions, mcParticlesPerMcCollision, mcParticles, Mass);
+    processGenMc(mcCollisions, mcParticlesPerMcCollision, mcParticles, rowsCommon.matchedCollisions);
   }
   PROCESS_SWITCH(HfDerivedDataCreatorDstarToD0Pi, processMcMlBkg, "Process MC with ML only for background", false);
 
@@ -398,13 +441,15 @@ struct HfDerivedDataCreatorDstarToD0Pi {
     rowsCommon.preProcessMcCollisions(mcCollisions, mcParticlesPerMcCollision, mcParticles);
     processCandidates<true, true, false, false>(collisions, candidatesMcMlAll, tracks, bcs);
     rowsCommon.processMcParticles(mcCollisions, mcParticlesPerMcCollision, mcParticles, Mass);
+    processGenMc(mcCollisions, mcParticlesPerMcCollision, mcParticles, rowsCommon.matchedCollisions);
   }
   PROCESS_SWITCH(HfDerivedDataCreatorDstarToD0Pi, processMcMlAll, "Process MC with ML", false);
 
   void processMcGenOnly(TypeMcCollisions const& mcCollisions,
                         MatchedGenCandidatesMc const& mcParticles)
   {
     rowsCommon.processMcParticles(mcCollisions, mcParticlesPerMcCollision, mcParticles, Mass);
+    processGenMc(mcCollisions, mcParticlesPerMcCollision, mcParticles, rowsCommon.matchedCollisions);
   }
   PROCESS_SWITCH(HfDerivedDataCreatorDstarToD0Pi, processMcGenOnly, "Process MC gen. only", false);
 };

PWGJE/DataModel/Jet.h

@@ -237,7 +237,7 @@ using CandidatesDstarMCD = o2::soa::Join<HfDstarBases, HfDstarPars, HfDstarParD0
 using JetTracksSubDstar = JTrackDstarSubs;
 using JetParticlesSubDstar = JMcParticleDstarSubs;
 using McCollisionsDstar = o2::soa::Join<HfDstarMcCollBases, JDstarMcCollisionIds>;
-using CandidatesDstarMCP = o2::soa::Join<HfDstarPBases, JDstarPIds>;
+using CandidatesDstarMCP = o2::soa::Join<HfDstarPBases, JDstarPIds, HfDstarGenMcs>;
 
 using CollisionsLc = o2::soa::Join<HfLcCollBases, JLcCollisionIds>;
 using CandidatesLcData = o2::soa::Join<HfLcBases, HfLcPars, HfLcParEs, HfLcSels, HfLcMls, JLcIds>;