Feat: finish mc processing for kinks

Author: ariedel-cern

PWGCF/Femto/Core/kinkBuilder.h

@@ -27,7 +27,6 @@
 #include "Common/Core/RecoDecay.h"
 
 #include "CommonConstants/MathConstants.h"
-#include "CommonConstants/PhysicsConstants.h"
 #include "Framework/AnalysisHelpers.h"
 #include "Framework/Configurable.h"
 
@@ -435,7 +434,6 @@ class KinkBuilder
     int64_t daughterIndex = 0;
 
     for (const auto& kink : kinks) {
-
       // compute mother kinematics before checking filters
       mKinkSelection.computeKinkMotherKinematics(kink);
       if (!mKinkSelection.checkFilters(kink)) {
@@ -444,15 +442,15 @@ class KinkBuilder
       // compute qa variables before applying selections
       mKinkSelection.computeQaVariables(kink, tracks);
       mKinkSelection.applySelections(kink, tracks);
-
       if (!mKinkSelection.passesAllRequiredSelections()) {
         continue;
       }
 
       collisionBuilder.template fillCollision<system>(collisionProducts, col);
       // cleaner, but without ITS pid: auto daughter = kink.template trackDaug_as<T7>();
       int64_t idx = kink.trackDaugId() - tracksWithItsPid.offset();
-      if (idx < 0 || idx > static_cast<int64_t>(tracksWithItsPid.size())) {
+      // check for valid index
+      if (idx < 0 || idx >= static_cast<int64_t>(tracksWithItsPid.size())) {
         return;
       }
       auto daughter = tracksWithItsPid.iteratorAt(idx);
@@ -466,6 +464,49 @@ class KinkBuilder
     }
   }
 
+  template <modes::System system, typename T1, typename T2, typename T3, typename T4, typename T5, typename T6, typename T7, typename T8, typename T9, typename T10, typename T11, typename T12, typename T13>
+  void fillMcKinks(T1 const& col, T2& collisionBuilder, T3& collisionProducts, T4 const& mcCols, T5& trackProducts, T6& kinkProducts, T7 const& kinks, T8 const& tracks, T9 const& tracksWithItsPid, T10& trackBuilder, T11 const& mcParticles, T12& mcBuilder, T13& mcProducts)
+  {
+
+    if (!mFillAnyTable) {
+      return;
+    }
+    int64_t daughterIndex = 0;
+    for (const auto& kink : kinks) {
+      // compute mother kinematics before checking filters
+      mKinkSelection.computeKinkMotherKinematics(kink);
+      if (!mKinkSelection.checkFilters(kink)) {
+        continue;
+      }
+      // compute qa variables before applying selections
+      mKinkSelection.computeQaVariables(kink, tracks);
+      mKinkSelection.applySelections(kink, tracks);
+      if (!mKinkSelection.passesAllRequiredSelections()) {
+        continue;
+      }
+
+      collisionBuilder.template fillMcCollision<system>(collisionProducts, col, mcCols, mcProducts, mcBuilder);
+
+      int64_t idx = kink.trackDaugId() - tracks.offset();
+      // check for valid index
+      if (idx < 0 || idx >= static_cast<int64_t>(tracks.size())) {
+        return;
+      }
+      auto daughter = tracks.iteratorAt(idx);
+      auto daughterWithItsPid = tracksWithItsPid.iteratorAt(idx);
+      daughterIndex = trackBuilder.template getDaughterIndex<system, modes::Track::kKinkDaughter>(col, collisionProducts, mcCols, daughter, daughterWithItsPid, trackProducts, mcParticles, mcBuilder, mcProducts);
+
+      if constexpr (modes::isEqual(kinkType, modes::Kink::kSigma)) {
+        fillSigma(collisionProducts, kinkProducts, kink, daughterIndex);
+        mcBuilder.template fillMcSigmaWithLabel<system>(col, mcCols, kink, daughter, mcParticles, mcProducts);
+      }
+      if constexpr (modes::isEqual(kinkType, modes::Kink::kSigmaPlus)) {
+        fillSigmaPlus(collisionProducts, kinkProducts, kink, daughterIndex);
+        mcBuilder.template fillMcSigmaPlusWithLabel<system>(col, mcCols, kink, daughter, mcParticles, mcProducts);
+      }
+    }
+  }
+
   template <typename T1, typename T2, typename T3>
   void fillSigma(T1& collisionProducts, T2& kinkProducts, T3 const& kink, int64_t daughterIndex)
   {

PWGCF/Femto/Core/kinkHistManager.h

@@ -52,6 +52,8 @@ struct McBuilderProducts : o2::framework::ProducesGroup {
   o2::framework::Produces<o2::aod::FTrackLabels> producedTrackLabels;
   o2::framework::Produces<o2::aod::FLambdaLabels> producedLambdaLabels;
   o2::framework::Produces<o2::aod::FK0shortLabels> producedK0shortLabels;
+  o2::framework::Produces<o2::aod::FSigmaLabels> producedSigmaLabels;
+  o2::framework::Produces<o2::aod::FSigmaPlusLabels> producedSigmaPlusLabels;
 };
 
 struct ConfMcTables : o2::framework::ConfigurableGroup {
@@ -65,6 +67,8 @@ struct ConfMcTables : o2::framework::ConfigurableGroup {
   o2::framework::Configurable<int> producedTrackLabels{"producedTrackLabels", -1, "Produce track labels (-1: auto; 0 off; 1 on)"};
   o2::framework::Configurable<int> producedLambdaLabels{"producedLambdaLabels", -1, "Produce lambda labels (-1: auto; 0 off; 1 on)"};
   o2::framework::Configurable<int> producedK0shortLabels{"producedK0shortLabels", -1, "Produce k0short labels (-1: auto; 0 off; 1 on)"};
+  o2::framework::Configurable<int> producedSigmaLabels{"producedSigmaLabels", -1, "Produce k0short labels (-1: auto; 0 off; 1 on)"};
+  o2::framework::Configurable<int> producedSigmaPlusLabels{"producedSigmaPlusLabels", -1, "Produce k0short labels (-1: auto; 0 off; 1 on)"};
 };
 
 class McBuilder
@@ -85,8 +89,10 @@ class McBuilder
 
     mProduceCollisionLabels = utils::enableTable("FColLabels", table.producedCollisionLabels.value, initContext);
     mProduceTrackLabels = utils::enableTable("FTrackLabels", table.producedTrackLabels.value, initContext);
-    mProduceLambdaLabels = utils::enableTable("FLambdaLabels", table.producedTrackLabels.value, initContext);
-    mProduceK0shortLabels = utils::enableTable("FK0shortLabels", table.producedTrackLabels.value, initContext);
+    mProduceLambdaLabels = utils::enableTable("FLambdaLabels", table.producedLambdaLabels.value, initContext);
+    mProduceK0shortLabels = utils::enableTable("FK0shortLabels", table.producedK0shortLabels.value, initContext);
+    mProduceSigmaLabels = utils::enableTable("FSigmaLabels", table.producedSigmaLabels.value, initContext);
+    mProduceSigmaPlusLabels = utils::enableTable("FSigmaPlusLabels", table.producedSigmaPlusLabels.value, initContext);
 
     if (mProduceMcCollisions || mProduceMcParticles || mProduceMcMothers || mProduceMcPartonicMothers || mProduceCollisionLabels || mProduceTrackLabels || mProduceLambdaLabels || mProduceK0shortLabels) {
       mFillAnyTable = true;
@@ -422,7 +428,195 @@ class McBuilder
         mK0shortToMcPartonicMap[k0shortIndex] = partonicRow;
       }
     }
-    mcProducts.producedTrackLabels(mcParticleRow, mothersRow, partonicRow);
+    mcProducts.producedK0shortLabels(mcParticleRow, mothersRow, partonicRow);
+  }
+
+  template <modes::System system, typename T1, typename T2, typename T3, typename T4, typename T5, typename T6>
+  void fillMcSigmaWithLabel(T1 const& col, T2 const& mcCols, T3 const& sigma, T4 const& sigmaDaughter, T5 const& mcParticles, T6& mcProducts)
+  {
+    if (!mProduceSigmaLabels) {
+      return;
+    }
+
+    if (!sigmaDaughter.has_mcParticle()) {
+      mcProducts.producedSigmaLabels(-1, -1, -1);
+      return;
+    }
+
+    auto mcKinkDaughter = sigmaDaughter.template mcParticle_as<T5>();
+    auto mcKinkDaughterMothers = mcKinkDaughter.template mothers_as<T5>();
+
+    if (!mcKinkDaughter.has_mothers() && !mcKinkDaughterMothers.empty()) {
+      mcProducts.producedSigmaLabels(-1, -1, -1);
+      return;
+    }
+
+    // we get the mc kink partilce via the mother of the kink daughter
+    auto mcParticle = mcKinkDaughterMothers.front();
+    auto mcCol = mcParticle.template mcCollision_as<T2>();
+
+    int64_t particleIndex = mcParticle.globalIndex();
+    int64_t sigmaIndex = sigma.globalIndex();
+
+    int64_t mcParticleRow = -1;
+
+    // MC PARTICLE
+    auto itP = mSigmaToMcParticleMap.find(particleIndex);
+    if (itP != mSigmaToMcParticleMap.end()) {
+      mcParticleRow = itP->second;
+    } else {
+      auto origin = this->getOrigin(col, mcCols, mcParticle);
+      int64_t mcColId = this->getMcColId<system>(mcCol, mcProducts);
+
+      mcProducts.producedMcParticles(
+        mcColId,
+        static_cast<aod::femtodatatypes::McOriginType>(origin),
+        mcParticle.pdgCode(),
+        mcParticle.pt() * utils::signum(mcParticle.pdgCode()),
+        mcParticle.eta(),
+        mcParticle.phi());
+
+      mcParticleRow = mcProducts.producedMcParticles.lastIndex();
+      mSigmaToMcParticleMap[particleIndex] = mcParticleRow;
+    }
+
+    // mothers (fill only if exists)
+    int64_t mothersRow = -1;
+    auto itM = mSigmaToMcMotherMap.find(sigmaIndex);
+
+    if (itM != mSigmaToMcMotherMap.end()) {
+      mothersRow = itM->second;
+    } else {
+
+      auto mothers = mcParticle.template mothers_as<T5>();
+      bool motherExists = mcParticle.has_mothers() && !mothers.empty();
+
+      if (motherExists) {
+        int motherPdg = mothers.front().pdgCode(); // PDG code is ALWAYS valid if the mother exists
+
+        mcProducts.producedMothers(motherPdg);
+        mothersRow = mcProducts.producedMothers.lastIndex();
+        mSigmaToMcMotherMap[sigmaIndex] = mothersRow;
+      }
+    }
+
+    // partonic mother (fill only if exists)
+    int64_t partonicRow = -1;
+    auto itPM = mSigmaToMcPartonicMap.find(sigmaIndex);
+
+    if (itPM != mSigmaToMcPartonicMap.end()) {
+      partonicRow = itPM->second;
+    } else {
+      int partIdx = -1;
+      if (mcParticle.has_mothers()) {
+        partIdx = this->findFirstPartonicMother(mcParticle, mcParticles);
+      }
+
+      bool partonicExists = (partIdx >= 0);
+
+      if (partonicExists) {
+        int partonicPdg = mcParticles.iteratorAt(partIdx).pdgCode();
+
+        mcProducts.producedPartonicMothers(partonicPdg);
+        partonicRow = mcProducts.producedPartonicMothers.lastIndex();
+        mSigmaToMcPartonicMap[sigmaIndex] = partonicRow;
+      }
+    }
+    mcProducts.producedSigmaLabels(mcParticleRow, mothersRow, partonicRow);
+  }
+
+  template <modes::System system, typename T1, typename T2, typename T3, typename T4, typename T5, typename T6>
+  void fillMcSigmaPlusWithLabel(T1 const& col, T2 const& mcCols, T3 const& sigmaPlus, T4 const& sigmaPlusDaughter, T5 const& mcParticles, T6& mcProducts)
+  {
+    if (!mProduceSigmaPlusLabels) {
+      return;
+    }
+
+    if (!sigmaPlusDaughter.has_mcParticle()) {
+      mcProducts.producedSigmaPlusLabels(-1, -1, -1);
+      return;
+    }
+
+    auto mcKinkDaughter = sigmaPlusDaughter.template mcParticle_as<T5>();
+    auto mcKinkDaughterMothers = mcKinkDaughter.template mothers_as<T5>();
+
+    if (!mcKinkDaughter.has_mothers() && !mcKinkDaughterMothers.empty()) {
+      mcProducts.producedSigmaPlusLabels(-1, -1, -1);
+      return;
+    }
+
+    // we get the mc kink partilce via the mother of the kink daughter
+    auto mcParticle = mcKinkDaughterMothers.front();
+    auto mcCol = mcParticle.template mcCollision_as<T2>();
+
+    int64_t particleIndex = mcParticle.globalIndex();
+    int64_t sigmaIndex = sigmaPlus.globalIndex();
+
+    int64_t mcParticleRow = -1;
+
+    // MC PARTICLE
+    auto itP = mSigmaPlusToMcParticleMap.find(particleIndex);
+    if (itP != mSigmaPlusToMcParticleMap.end()) {
+      mcParticleRow = itP->second;
+    } else {
+      auto origin = this->getOrigin(col, mcCols, mcParticle);
+      int64_t mcColId = this->getMcColId<system>(mcCol, mcProducts);
+
+      mcProducts.producedMcParticles(
+        mcColId,
+        static_cast<aod::femtodatatypes::McOriginType>(origin),
+        mcParticle.pdgCode(),
+        mcParticle.pt() * utils::signum(mcParticle.pdgCode()),
+        mcParticle.eta(),
+        mcParticle.phi());
+
+      mcParticleRow = mcProducts.producedMcParticles.lastIndex();
+      mSigmaPlusToMcParticleMap[particleIndex] = mcParticleRow;
+    }
+
+    // mothers (fill only if exists)
+    int64_t mothersRow = -1;
+    auto itM = mSigmaPlusToMcMotherMap.find(sigmaIndex);
+
+    if (itM != mSigmaPlusToMcMotherMap.end()) {
+      mothersRow = itM->second;
+    } else {
+
+      auto mothers = mcParticle.template mothers_as<T5>();
+      bool motherExists = mcParticle.has_mothers() && !mothers.empty();
+
+      if (motherExists) {
+        int motherPdg = mothers.front().pdgCode(); // PDG code is ALWAYS valid if the mother exists
+
+        mcProducts.producedMothers(motherPdg);
+        mothersRow = mcProducts.producedMothers.lastIndex();
+        mSigmaPlusToMcMotherMap[sigmaIndex] = mothersRow;
+      }
+    }
+
+    // partonic mother (fill only if exists)
+    int64_t partonicRow = -1;
+    auto itPM = mSigmaPlusToMcPartonicMap.find(sigmaIndex);
+
+    if (itPM != mSigmaPlusToMcPartonicMap.end()) {
+      partonicRow = itPM->second;
+    } else {
+      int partIdx = -1;
+      if (mcParticle.has_mothers()) {
+        partIdx = this->findFirstPartonicMother(mcParticle, mcParticles);
+      }
+
+      bool partonicExists = (partIdx >= 0);
+
+      if (partonicExists) {
+        int partonicPdg = mcParticles.iteratorAt(partIdx).pdgCode();
+
+        mcProducts.producedPartonicMothers(partonicPdg);
+        partonicRow = mcProducts.producedPartonicMothers.lastIndex();
+        mSigmaPlusToMcPartonicMap[sigmaIndex] = partonicRow;
+      }
+    }
+    mcProducts.producedSigmaPlusLabels(mcParticleRow, mothersRow, partonicRow);
   }
 
   template <typename TParticle, typename TContainer>
@@ -488,8 +682,18 @@ class McBuilder
     mK0shortToMcParticleMap.clear();
     mK0shortToMcMotherMap.clear();
     mK0shortToMcPartonicMap.clear();
+
+    mSigmaToMcParticleMap.clear();
+    mSigmaToMcMotherMap.clear();
+    mSigmaToMcPartonicMap.clear();
+
+    mSigmaPlusToMcParticleMap.clear();
+    mSigmaPlusToMcMotherMap.clear();
+    mSigmaPlusToMcPartonicMap.clear();
   }
 
+  bool fillAnyTable() const { return mFillAnyTable; }
+
  private:
   bool mPassThrough = false;
   bool mFillAnyTable = false;
@@ -502,6 +706,8 @@ class McBuilder
   bool mProduceTrackLabels = false;
   bool mProduceLambdaLabels = false;
   bool mProduceK0shortLabels = false;
+  bool mProduceSigmaLabels = false;
+  bool mProduceSigmaPlusLabels = false;
 
   std::unordered_map<int64_t, int64_t> mCollisionMap;
 
@@ -516,6 +722,14 @@ class McBuilder
   std::unordered_map<int64_t, int64_t> mK0shortToMcParticleMap;
   std::unordered_map<int64_t, int64_t> mK0shortToMcMotherMap;
   std::unordered_map<int64_t, int64_t> mK0shortToMcPartonicMap;
+
+  std::unordered_map<int64_t, int64_t> mSigmaToMcParticleMap;
+  std::unordered_map<int64_t, int64_t> mSigmaToMcMotherMap;
+  std::unordered_map<int64_t, int64_t> mSigmaToMcPartonicMap;
+
+  std::unordered_map<int64_t, int64_t> mSigmaPlusToMcParticleMap;
+  std::unordered_map<int64_t, int64_t> mSigmaPlusToMcMotherMap;
+  std::unordered_map<int64_t, int64_t> mSigmaPlusToMcPartonicMap;
 };
 
 } // namespace mcbuilder

PWGCF/Femto/Core/mcBuilder.h

@@ -269,10 +269,12 @@ struct ConfTrackMcBinning : o2::framework::ConfigurableGroup {
 constexpr const char PrefixTrackMcBinning1[] = "TrackMcBinning1";
 constexpr const char PrefixV0PosDauMcBinning[] = "V0PosDauMcBinning";
 constexpr const char PrefixV0NegDauMcBinning[] = "V0NegDauMcBinning";
+constexpr const char PrefixKinkChaDauMcBinning[] = "KinkChaDauMcBinning";
 
 using ConfTrackMcBinning1 = ConfTrackMcBinning<PrefixTrackMcBinning1>;
 using ConfV0PosDauMcBinning = ConfTrackMcBinning<PrefixV0PosDauMcBinning>;
 using ConfV0NegDauMcBinning = ConfTrackMcBinning<PrefixV0NegDauMcBinning>;
+using ConfKinkChaDauMcBinning = ConfTrackMcBinning<PrefixKinkChaDauMcBinning>;
 
 // must be in sync with enum TrackVariables
 // the enum gives the correct index in the array

PWGCF/Femto/Core/trackHistManager.h

@@ -177,9 +177,9 @@ constexpr std::array<histmanager::HistInfo<V0Hist>, kV0HistLast> HistTable = {
    {kK0shortMassVsAntiLambdaMass, o2::framework::kTH2F, "hK0shortMassVsAntiLambdaMass", "K^{0}_{S} mass vs #bar{#Lambda} mass; m_{#pi^{+}#pi^{-}} (GeV/#it{c}^{2}); m_{#bar{p}#pi^{+}} (GeV/#it{c}^{2})"},
    {kLambdaMassVsAntiLambdaMass, o2::framework::kTH2F, "hLambdaMassVsAntiLambdaMass", "#Lambda mass vs #bar{#Lambda}; m_{p#pi^{-}} (GeV/#it{c}^{2}); m_{#bar{p}#pi^{+}} (GeV/#it{c}^{2})"},
    {kOrigin, o2::framework::kTH1F, "hOrigin", "Status Codes (=Origin); Status Code; Entries"},
-   {kPdg, o2::framework::kTH1F, "hPdg", "PDG Codes of selected tracks; PDG Code; Entries"},
-   {kPdgMother, o2::framework::kTH1F, "hPdgMother", "PDG Codes of mother of selected tracks; PDG Code; Entries"},
-   {kPdgPartonicMother, o2::framework::kTH1F, "hPdgPartonicMother", "PDG Codes of partonic mother selected tracks; PDG Code; Entries"},
+   {kPdg, o2::framework::kTH1F, "hPdg", "PDG Codes of reconstructed v0; PDG Code; Entries"},
+   {kPdgMother, o2::framework::kTH1F, "hPdgMother", "PDG Codes of mother of reconstructed v0; PDG Code; Entries"},
+   {kPdgPartonicMother, o2::framework::kTH1F, "hPdgPartonicMother", "PDG Codes of partonic mother of reconstructed v0; PDG Code; Entries"},
    {kTruePt, o2::framework::kTH1F, "hTruePt", "True transverse momentum; p_{T} (GeV/#it{c}); Entries"},
    {kTrueEta, o2::framework::kTH1F, "hTrueEta", "True pseudorapdity; #eta; Entries"},
    {kTruePhi, o2::framework::kTH1F, "hTruePhi", "True azimuthal angle; #varphi; Entries"},
@@ -407,6 +407,7 @@ class V0HistManager
   }
 
  private:
+
   // for qa
   template <typename T1>
   void enableOptionalHistograms(T1 const& V0ConfBinningQa)
@@ -418,7 +419,7 @@ class V0HistManager
   template <typename T1, typename T2>
   void enableOptionalHistograms(T1 const& V0ConfBinningQa, T2 const& V0ConfBinningMc)
   {
-    mPlot2d = V0ConfBinningQa.plot2d.value;
+    this->enableOptionalHistograms(V0ConfBinningQa);
 
     mPlotOrigins = V0ConfBinningMc.plotOrigins.value;
     mPlotNSecondaries = V0ConfBinningMc.pdgCodesForMothersOfSecondary.value.size();