[PWGLF] Align K1 MC parent truth and efficiency response

Author: Bong-Hwi Lim

PWGLF/Core/K1AnalysisMicroCore.h

@@ -43,6 +43,7 @@
 #include <limits>
 #include <set>
 #include <type_traits>
+#include <unordered_map>
 #include <utility>
 #include <vector>
 
@@ -64,6 +65,16 @@ enum class K1TruthChannel {
   KStarPi = 2
 };
 
+struct K1ParentTruth {
+  double pt = std::numeric_limits<double>::quiet_NaN();
+  double y = std::numeric_limits<double>::quiet_NaN();
+  int pdgCode = 0;
+  K1TruthChannel channel = K1TruthChannel::None;
+};
+using K1ParentTruthMap = std::unordered_map<int64_t, K1ParentTruth>;
+
+inline K1TruthChannel classifyGeneratedK1(int charge, int daughter1, int daughter2);
+
 // The value is the species index of the PID configuration in ResoAnalysisSelectionCore.
 enum class Species : int {
   Pion = 0,
@@ -217,6 +228,49 @@ K1TruthChannel classifyK1Truth(const Track& pion1, const Track& pion2, const Kao
   return K1TruthChannel::None;
 }
 
+template <typename Kaon, typename Pion>
+int k1MotherId(K1TruthChannel channel, const Kaon& kaon, const Pion& pion1, const Pion& pion2)
+{
+  if (channel == K1TruthChannel::RhoK) {
+    return kaon.motherId();
+  }
+  if (channel == K1TruthChannel::KStarPi) {
+    return std::abs(pion1.motherPDG()) == o2::constants::physics::Pdg::kK1_1270Plus ? pion1.motherId() : pion2.motherId();
+  }
+  return -1;
+}
+
+template <typename ParentsType>
+K1ParentTruthMap makeK1ParentTruthMap(const ParentsType& parents)
+{
+  K1ParentTruthMap result;
+  for (const auto& part : parents) {
+    if (std::abs(part.pdgCode()) != o2::constants::physics::Pdg::kK1_1270Plus || part.originalMcParticleId() < 0) {
+      continue;
+    }
+    const auto id = static_cast<int64_t>(part.originalMcParticleId());
+    const int charge = part.pdgCode() > 0 ? 1 : -1;
+    result.emplace(id, K1ParentTruth{part.pt(), part.y(), part.pdgCode(), classifyGeneratedK1(charge, part.daughterPDG1(), part.daughterPDG2())});
+  }
+  return result;
+}
+
+inline const K1ParentTruth* findK1ParentTruth(const K1ParentTruthMap& parents, int64_t motherId,
+                                              K1TruthChannel channel, int charge)
+{
+  const auto it = parents.find(motherId);
+  if (it == parents.end()) {
+    return nullptr;
+  }
+  const auto& parent = it->second;
+  const int expectedPdg = charge * o2::constants::physics::Pdg::kK1_1270Plus;
+  if (channel == K1TruthChannel::None || parent.channel != channel || parent.pdgCode != expectedPdg ||
+      !std::isfinite(parent.pt) || parent.pt < 0. || !std::isfinite(parent.y)) {
+    return nullptr;
+  }
+  return &parent;
+}
+
 // Immediate decay channel of a generated K1 from the PDG codes of its two daughters.
 inline K1TruthChannel classifyGeneratedK1(int charge, int daughter1, int daughter2)
 {
@@ -635,7 +689,8 @@ class K1AnalysisMicroCore
       const int charge = part.pdgCode() > 0 ? 1 : -1;
       const K1TruthChannel channel = classifyGeneratedK1(charge, part.daughterPDG1(), part.daughterPDG2());
       histos.fill(HIST("CutFlow/generated"), 0, static_cast<int>(channel));
-      if (part.y() < mCandidateCuts.cK1MinRap || part.y() > mCandidateCuts.cK1MaxRap) {
+      if (!std::isfinite(part.pt()) || part.pt() < 0. || !std::isfinite(part.y()) ||
+          part.y() < mCandidateCuts.cK1MinRap || part.y() > mCandidateCuts.cK1MaxRap) {
         continue;
       }
       histos.fill(HIST("CutFlow/generated"), 1, static_cast<int>(channel));

PWGLF/DataModel/LFK1MlTables.h

@@ -69,8 +69,8 @@ DECLARE_SOA_COLUMN(K1TruthStatus, k1TruthStatus, uint8_t);          //! 0 data,
 DECLARE_SOA_COLUMN(K1TruthChannel, k1TruthChannel, uint8_t);        //! 0 none, 1 rho K, 2 K* pi
 DECLARE_SOA_COLUMN(K1MotherPdg, k1MotherPdg, int32_t);
 DECLARE_SOA_COLUMN(K1MotherId, k1MotherId, int64_t);
-DECLARE_SOA_COLUMN(K1GeneratedPt, k1GeneratedPt, float); //! NaN in K1MlTruth: not filled for reconstructed candidates
-DECLARE_SOA_COLUMN(K1GeneratedY, k1GeneratedY, float);   //! NaN in K1MlTruth: not filled for reconstructed candidates
+DECLARE_SOA_COLUMN(K1GeneratedPt, k1GeneratedPt, float); //! Parent-table pT for a validated reconstructed MC match; NaN otherwise
+DECLARE_SOA_COLUMN(K1GeneratedY, k1GeneratedY, float);   //! Parent-table y for a validated reconstructed MC match; NaN otherwise
 DECLARE_SOA_COLUMN(K1OriginalMcParticleId, k1OriginalMcParticleId, int64_t);
 DECLARE_SOA_COLUMN(K1DaughterPdg1, k1DaughterPdg1, int32_t);
 DECLARE_SOA_COLUMN(K1DaughterPdg2, k1DaughterPdg2, int32_t);

PWGLF/Tasks/Resonances/k1AnalysisMicro.cxx

@@ -52,6 +52,7 @@ struct HistogramOptions : ConfigurableGroup {
   Configurable<int> cNbinsDiv{"cNbinsDiv", 1, "Integer to divide the number of bins"};
   Configurable<bool> additionalQAplots{"additionalQAplots", true, "Additional QA plots"};
   Configurable<int> cfgTruthDebug{"cfgTruthDebug", 0, "Maximum logged matched candidates per truth channel"};
+  Configurable<bool> cUseTruthRapidity{"cUseTruthRapidity", true, "Require generated K1 rapidity in the configured window for MC efficiency numerators"};
 };
 
 enum BinAnti : unsigned int {
@@ -231,12 +232,15 @@ struct K1AnalysisMicro {
       AxisSpec channelAxis = {3, -0.5, 2.5, "0: non-K1, 1: rho K, 2: K* pi"};
       histos.add("MCReco/collisions", "Selected reconstructed MC collisions", HistType::kTH1D, {{1, 0, 1}});
       histos.add("MCReco/microTracks", "Input micro tracks in selected MC collisions", HistType::kTH1D, {{1, 0, 1}});
-      histos.add("MCReco/channel", "All selected pi-pi-K combinations by truth channel", HistType::kTH1D, {channelAxis});
+      // These channel and pT histograms are pre-truth-fiducial QA; MCReco/response is the validated numerator.
+      histos.add("MCReco/channel", "Pre-truth-fiducial QA: selected pi-pi-K combinations by truth channel", HistType::kTH1D, {channelAxis});
       histos.add("MCReco/mass", "Reconstructed mass by truth channel", HistType::kTH2D, {channelAxis, invMassAxisReso});
-      histos.add("MCReco/pt", "Reconstructed pT by truth channel", HistType::kTH2D, {channelAxis, ptAxis});
+      histos.add("MCReco/pt", "Pre-truth-fiducial QA: reconstructed pT by truth channel", HistType::kTH2D, {channelAxis, ptAxis});
       histos.add("MCReco/piPiMass", "pi-pi mass by truth channel", HistType::kTH2D, {channelAxis, invMassAxisRho});
       histos.add("MCReco/pi1KMass", "First pion-kaon mass by truth channel", HistType::kTH2D, {channelAxis, invMassAxisK892});
       histos.add("MCReco/pi2KMass", "Second pion-kaon mass by truth channel", HistType::kTH2D, {channelAxis, invMassAxisK892});
+      histos.add("MCReco/response", "Selected MC candidates;reconstructed pT;generated pT;channel;charge;centrality", HistType::kTHnSparseD,
+                 {ptAxis, ptAxis, channelAxis, AxisSpec{2, -1.5, 1.5, "K1 charge"}, centAxis});
       histos.add("k1invmass_MC", "Invariant mass of K1(1270)", HistType::kTH1F, {invMassAxisReso});
       histos.add("k1invmass_MC_noK1", "Invariant mass of K1(1270)", HistType::kTH1F, {invMassAxisReso});
 
@@ -272,6 +276,8 @@ struct K1AnalysisMicro {
       AxisSpec channelAxis = {3, -0.5, 2.5, "0: other/unresolved, 1: rho K, 2: K* pi"};
       histos.add("MCGen/chargeChannel", "K1 parents in selected reconstructed events, inside the K1 rapidity window", HistType::kTH2D, {{2, -1.5, 1.5, "K1 charge"}, channelAxis});
       histos.add("MCGen/ptChannel", "Generated K1 pT by immediate decay channel", HistType::kTH2D, {channelAxis, ptAxis});
+      histos.add("MCGen/response", "Generated K1 parents;generated pT;channel;charge;centrality", HistType::kTHnSparseD,
+                 {ptAxis, channelAxis, AxisSpec{2, -1.5, 1.5, "K1 charge"}, centAxis});
     }
   }
 
@@ -380,7 +386,8 @@ struct K1AnalysisMicro {
   // (or one mixed pair of collisions). The selection itself is the shared K1 core.
   // dTracks1: bachelor kaons, dTracks2: pions.
   template <bool IsMC, bool IsMix, bool IsResoMicrotrack, typename CollisionType, typename TracksType>
-  void fillHistograms(const CollisionType& collision, const TracksType& dTracks1, const TracksType& dTracks2)
+  void fillHistograms(const CollisionType& collision, const TracksType& dTracks1, const TracksType& dTracks2,
+                      const K1ParentTruthMap* parentTruth = nullptr)
   {
     const bool fillQA = !IsMix && histogramOptions.additionalQAplots;
     const auto multiplicity = collision.cent();
@@ -474,9 +481,17 @@ struct K1AnalysisMicro {
                pion1.motherPDG(), pion2.motherPDG(), kaon.motherPDG(),
                pion1.siblingIds()[0], pion1.siblingIds()[1], pion2.siblingIds()[0], pion2.siblingIds()[1], kaon.siblingIds()[0], kaon.siblingIds()[1]);
         }
-        typeK1 = kaon.sign() > 0 ? BinType::kK1P_Rec : BinType::kK1N_Rec;
-        histos.fill(HIST("hInvmass_K1"), typeNormal, typeK1, multiplicity, c.k1.Pt(), c.k1.M());
-        histos.fill(HIST("k1invmass_MC"), c.k1.M());
+        const int motherId = k1MotherId(channel, kaon, pion1, pion2);
+        const int charge = kaon.pdgCode() > 0 ? 1 : -1;
+        const auto* parent = parentTruth ? findK1ParentTruth(*parentTruth, motherId, channel, charge) : nullptr;
+        const bool truthRapidity = c.isUnlikeSign && parent && (!histogramOptions.cUseTruthRapidity || (parent->y >= candidateCuts.cK1MinRap && parent->y <= candidateCuts.cK1MaxRap));
+        if (truthRapidity) {
+          typeK1 = kaon.sign() > 0 ? BinType::kK1P_Rec : BinType::kK1N_Rec;
+          histos.fill(HIST("hInvmass_K1"), typeNormal, typeK1, multiplicity, c.k1.Pt(), c.k1.M());
+          histos.fill(HIST("k1invmass_MC"), c.k1.M());
+          const int parentCharge = parent->pdgCode > 0 ? 1 : -1;
+          histos.fill(HIST("MCReco/response"), c.k1.Pt(), parent->pt, channelBin, parentCharge, multiplicity);
+        }
         histos.fill(HIST("QAMC/K1OA"), c.angle);
         histos.fill(HIST("QAMC/K1PairAsym"), c.pairAsym);
         histos.fill(HIST("QAMC/hInvmassK892_Rho"), c.mass13, c.secondary.M());
@@ -515,17 +530,18 @@ struct K1AnalysisMicro {
   PROCESS_SWITCH(K1AnalysisMicro, processResoMicroTracks, "Process ResoMicroTracks", true);
 
   void processMC(ResoMCCols::iterator const& collision,
-                 ResoMCTracks const& resotracks)
+                 ResoMCTracks const& resotracks, ResoMCParents const& resoParents)
   {
     if (!core.passesEventCuts(collision) || !core.passesMCEventCuts(collision)) {
       return;
     }
     histos.fill(HIST("MCReco/collisions"), 0.5);
-    fillHistograms<true, false, false>(collision, resotracks, resotracks);
+    const auto parentTruth = makeK1ParentTruthMap(resoParents);
+    fillHistograms<true, false, false>(collision, resotracks, resotracks, &parentTruth);
   }
   PROCESS_SWITCH(K1AnalysisMicro, processMC, "Process Event for MC", false);
 
-  void processMCMicro(ResoMCCols::iterator const& collision, ResoMCMicroTracks const& tracks)
+  void processMCMicro(ResoMCCols::iterator const& collision, ResoMCMicroTracks const& tracks, ResoMCParents const& resoParents)
   {
     // The modular producer already selected these reconstructed collisions.
     // Apply precisely the same reconstruction loop as the data baseline.
@@ -534,7 +550,8 @@ struct K1AnalysisMicro {
     }
     histos.fill(HIST("MCReco/collisions"), 0.5);
     histos.fill(HIST("MCReco/microTracks"), 0.5, tracks.size());
-    fillHistograms<true, false, true>(collision, tracks, tracks);
+    const auto parentTruth = makeK1ParentTruthMap(resoParents);
+    fillHistograms<true, false, true>(collision, tracks, tracks, &parentTruth);
   }
   PROCESS_SWITCH(K1AnalysisMicro, processMCMicro, "Process reconstructed MC with micro v001 tables", false);
 
@@ -548,6 +565,9 @@ struct K1AnalysisMicro {
       const int charge = part.pdgCode() > 0 ? 1 : -1;
       histos.fill(HIST("MCGen/chargeChannel"), charge, static_cast<int>(channel));
       histos.fill(HIST("MCGen/ptChannel"), static_cast<int>(channel), part.pt());
+      if (channel != K1TruthChannel::None) {
+        histos.fill(HIST("MCGen/response"), part.pt(), static_cast<int>(channel), charge, collision.cent());
+      }
     });
   }
   PROCESS_SWITCH(K1AnalysisMicro, processMCTrue, "Process generated K1 in selected events with v001 parents", false);

PWGLF/Tasks/Resonances/k1TrainingTable.cxx

@@ -182,7 +182,7 @@ struct K1TrainingTable {
 
   template <bool IsMC, typename Collision, typename Kaon, typename Pion>
   void writeK1MlCandidate(Collision const& collision, Kaon const& kaon, Pion const& samePion, Pion const& oppPion,
-                          K1TruthChannel channel, uint16_t passBits)
+                          K1TruthChannel channel, uint16_t passBits, const K1ParentTruthMap* parentTruth = nullptr)
   {
     using o2::analysis::k1ml::BuildStatus;
     const auto canonical = o2::analysis::k1ml::canonicalizeUS(o2::analysis::k1ml::makeTrackSnapshot(kaon), o2::analysis::k1ml::makeTrackSnapshot(samePion), o2::analysis::k1ml::makeTrackSnapshot(oppPion));
@@ -215,14 +215,17 @@ struct K1TrainingTable {
     if constexpr (IsMC) {
       const bool matched = channel != K1TruthChannel::None;
       int motherId = -1;
+      const K1ParentTruth* parent = nullptr;
       if (matched) {
-        motherId = channel == K1TruthChannel::RhoK ? kaon.motherId() : std::abs(samePion.motherPDG()) == Pdg::kK1_1270Plus ? samePion.motherId()
-                                                                                                                           : oppPion.motherId();
+        motherId = k1MotherId(channel, kaon, samePion, oppPion);
+        const int charge = kaon.pdgCode() > 0 ? 1 : -1;
+        parent = parentTruth ? findK1ParentTruth(*parentTruth, motherId, channel, charge) : nullptr;
       }
       // Immediate-mother information alone does not prove a physical UID.
       k1MlTruth(row, matched ? uint8_t{1} : uint8_t{2}, static_cast<uint8_t>(channel),
-                matched ? kaon.sign() * Pdg::kK1_1270Plus : 0, static_cast<int64_t>(motherId),
-                std::numeric_limits<float>::quiet_NaN(), std::numeric_limits<float>::quiet_NaN());
+                parent ? parent->pdgCode : (matched ? kaon.sign() * Pdg::kK1_1270Plus : 0), static_cast<int64_t>(motherId),
+                parent ? static_cast<float>(parent->pt) : std::numeric_limits<float>::quiet_NaN(),
+                parent ? static_cast<float>(parent->y) : std::numeric_limits<float>::quiet_NaN());
     } else {
       k1MlTruth(row, uint8_t{0}, uint8_t{0}, 0, int64_t{-1},
                 std::numeric_limits<float>::quiet_NaN(), std::numeric_limits<float>::quiet_NaN());
@@ -244,7 +247,7 @@ struct K1TrainingTable {
   }
   PROCESS_SWITCH(K1TrainingTable, processResoMicroTracks, "Write K1 candidates from data micro v001 tables", true);
 
-  void processMCMicro(ResoMCCols::iterator const& collision, ResoMCMicroTracks const& tracks)
+  void processMCMicro(ResoMCCols::iterator const& collision, ResoMCMicroTracks const& tracks, ResoMCParents const& resoParents)
   {
     // The modular producer already selected these reconstructed collisions.
     if (!core.passesEventCuts(collision) || !core.passesMCEventCuts(collision)) {
@@ -253,10 +256,11 @@ struct K1TrainingTable {
     histos.fill(HIST("MCReco/collisions"), 0.5);
     histos.fill(HIST("MCReco/microTracks"), 0.5, tracks.size());
     writeK1MlEvent(collision);
+    const auto parentTruth = makeK1ParentTruthMap(resoParents);
     core.forEachCandidate<true, false, true>(histos, collision, tracks, tracks, false, nullptr, nullptr,
-                                             [this](auto const& coll, auto const& kaon, auto const& samePion, auto const& oppPion,
-                                                    K1TruthChannel channel, uint16_t passBits) {
-                                               writeK1MlCandidate<true>(coll, kaon, samePion, oppPion, channel, passBits);
+                                             [this, &parentTruth](auto const& coll, auto const& kaon, auto const& samePion, auto const& oppPion,
+                                                                  K1TruthChannel channel, uint16_t passBits) {
+                                               writeK1MlCandidate<true>(coll, kaon, samePion, oppPion, channel, passBits, &parentTruth);
                                              });
   }
   PROCESS_SWITCH(K1TrainingTable, processMCMicro, "Write K1 candidates with truth from reconstructed MC micro v001 tables", false);