[PWGHF] Add MC processing support to UPC Lc task (#17876)

Author: Rrantu

PWGHF/D2H/Tasks/taskUpcLc.cxx

@@ -47,6 +47,7 @@
 
 #include <array>
 #include <cmath>
+#include <cstdint>
 #include <numeric>
 #include <string>
 #include <vector> // std::vector
@@ -79,6 +80,9 @@ DECLARE_SOA_COLUMN(AmpFT0A, ampFT0A, float);
 DECLARE_SOA_COLUMN(AmpFT0C, ampFT0C, float);
 DECLARE_SOA_COLUMN(ZdcTimeZNA, zdcTimeZNA, float);
 DECLARE_SOA_COLUMN(ZdcTimeZNC, zdcTimeZNC, float);
+DECLARE_SOA_COLUMN(FlagMcMatchRec, flagMcMatchRec, int8_t);
+DECLARE_SOA_COLUMN(OriginMcRec, originMcRec, int8_t);
+DECLARE_SOA_COLUMN(PtBhadMotherPart, ptBhadMotherPart, float);
 } // namespace full
 DECLARE_SOA_TABLE(HfUpcQa, "AOD", "HFUPCQA",
                   full::PvContributors,
@@ -103,12 +107,33 @@ DECLARE_SOA_TABLE(HfUpcLcInfos, "AOD", "HFUPCLCINFOS",
                   full::Chi2PCA,
                   full::DecayLength,
                   full::Cpa);
+DECLARE_SOA_TABLE(HfUpcLcMcBdtInfos, "AOD", "HFUPCLCMCBDT",
+                  full::M,
+                  full::Pt,
+                  full::BkgScore,
+                  full::FlagMcMatchRec,
+                  full::OriginMcRec,
+                  full::PtBhadMotherPart);
+DECLARE_SOA_TABLE(HfUpcLcMcInfos, "AOD", "HFUPCLCMCINFO",
+                  full::M,
+                  full::Pt,
+                  full::PtProng0,
+                  full::PtProng1,
+                  full::PtProng2,
+                  full::Chi2PCA,
+                  full::DecayLength,
+                  full::Cpa,
+                  full::FlagMcMatchRec,
+                  full::OriginMcRec,
+                  full::PtBhadMotherPart);
 } // namespace o2::aod
 
 /// Λc± → p± K∓ π± analysis task
 struct HfTaskUpcLc {
   Produces<o2::aod::HfUpcLcBdtInfos> rowCandUpcBdt;
   Produces<o2::aod::HfUpcLcInfos> rowCandUpc;
+  Produces<o2::aod::HfUpcLcMcBdtInfos> rowCandUpcMcBdt;
+  Produces<o2::aod::HfUpcLcMcInfos> rowCandUpcMc;
   Produces<o2::aod::HfUpcQa> rowUpcQa;
 
   Configurable<int> selectionFlagLc{"selectionFlagLc", 1, "Selection Flag for Lc"};
@@ -133,6 +158,8 @@ struct HfTaskUpcLc {
 
   using LcCandidates = soa::Filtered<soa::Join<aod::HfCand3Prong, aod::HfSelLc>>;
   using LcCandidatesMl = soa::Filtered<soa::Join<aod::HfCand3Prong, aod::HfSelLc, aod::HfMlLcToPKPi>>;
+  using LcCandidatesMc = soa::Filtered<soa::Join<aod::HfCand3Prong, aod::HfSelLc, aod::HfCand3ProngMcRec>>;
+  using LcCandidatesMlMc = soa::Filtered<soa::Join<aod::HfCand3Prong, aod::HfSelLc, aod::HfMlLcToPKPi, aod::HfCand3ProngMcRec>>;
 
   Filter filterSelectCandidates = aod::hf_sel_candidate_lc::isSelLcToPKPi >= selectionFlagLc || aod::hf_sel_candidate_lc::isSelLcToPiKP >= selectionFlagLc;
   Preslice<aod::HfCand3Prong> candLcPerCollision = aod::hf_cand::collisionId;
@@ -149,7 +176,7 @@ struct HfTaskUpcLc {
 
   void init(InitContext&)
   {
-    const std::array<bool, 2> doprocess{doprocessDataWithMlWithUpc, doprocessDataStdWithUpc};
+    const std::array<bool, 4> doprocess{doprocessDataWithMlWithUpc, doprocessDataStdWithUpc, doprocessMcWithMlWithUpc, doprocessMcStdWithUpc};
     if ((std::accumulate(doprocess.begin(), doprocess.end(), 0)) != 1) {
       LOGP(fatal, "no or more than one process function enabled! Please check your configuration!");
     }
@@ -178,13 +205,13 @@ struct HfTaskUpcLc {
     return o2::hf_centrality::getCentralityColl<Coll>(collision);
   }
 
-  template <bool FillMl, typename CollType, typename CandType, typename BCsType>
-  void runAnalysisPerCollisionDataWithUpc(CollType const& collisions,
-                                          CandType const& candidates,
-                                          BCsType const& bcs,
-                                          aod::FT0s const& ft0s,
-                                          aod::FV0As const& fv0as,
-                                          aod::FDDs const& fdds
+  template <bool FillMl, bool IsMc, typename CollType, typename CandType, typename BCsType>
+  void runAnalysisPerCollisionWithUpc(CollType const& collisions,
+                                      CandType const& candidates,
+                                      BCsType const& bcs,
+                                      aod::FT0s const& ft0s,
+                                      aod::FV0As const& fv0as,
+                                      aod::FDDs const& fdds
 
   )
   {
@@ -244,7 +271,23 @@ struct HfTaskUpcLc {
         }
         registry.fill(HIST("Data/hUpcGapAfterSelection"), static_cast<int>(gap));
       }
-      const bool ignoreZdcTime = (zdcTimeThreshold < 0.f);
+      if constexpr (!IsMc) {
+        if (!hasZdc) {
+          continue;
+        }
+      }
+      if constexpr (IsMc) {
+        if (!hasZdc && fillHistQa) {
+          registry.fill(HIST("Data/fitInfo/ampFT0A_vs_ampFT0C"), fitInfo.ampFT0A, fitInfo.ampFT0C);
+        }
+        if (!hasZdc) {
+          registry.fill(HIST("Data/hUpcGapAfterSelection"), static_cast<int>(gap));
+        }
+      }
+      bool ignoreZdcTime = (zdcTimeThreshold < 0.f);
+      if constexpr (IsMc) {
+        ignoreZdcTime = ignoreZdcTime || !hasZdc;
+      }
       const auto multNTracksPV = collision.multNTracksPV();
       const auto posZ = collision.posZ();
       if (gap == o2::aod::sgselector::TrueGap::SingleGapA && (ignoreZdcTime || (std::abs(zdcTimeZNA) > zdcTimeThreshold && std::abs(zdcTimeZNC) < zdcTimeThreshold))) {
@@ -289,6 +332,18 @@ struct HfTaskUpcLc {
 
         auto fillTHnData = [&](bool isPKPi) {
           const auto massLc = isPKPi ? HfHelper::invMassLcToPKPi(candidate) : HfHelper::invMassLcToPiKP(candidate);
+          if constexpr (IsMc) {
+            if constexpr (FillMl) {
+              const auto& mlProb = isPKPi ? candidate.mlProbLcToPKPi() : candidate.mlProbLcToPiKP();
+              if (mlProb.size() == NumberOfMlClasses) {
+                outputBkg = mlProb[MlClassBackground]; /// bkg score
+              }
+              rowCandUpcMcBdt(massLc, pt, outputBkg, candidate.flagMcMatchRec(), candidate.originMcRec(), candidate.ptBhadMotherPart());
+            } else {
+              rowCandUpcMc(massLc, pt, ptProng0, ptProng1, ptProng2, chi2PCA, decayLength, cpa, candidate.flagMcMatchRec(), candidate.originMcRec(), candidate.ptBhadMotherPart());
+            }
+            return;
+          }
           if constexpr (FillMl) {
             const auto& mlProb = isPKPi ? candidate.mlProbLcToPKPi() : candidate.mlProbLcToPiKP();
             if (mlProb.size() == NumberOfMlClasses) {
@@ -321,7 +376,7 @@ struct HfTaskUpcLc {
                                 aod::FDDs const& fdds,
                                 aod::Zdcs const& /*zdcs*/)
   {
-    runAnalysisPerCollisionDataWithUpc<true>(collisions, selectedLcCandidatesMl, bcs, ft0s, fv0as, fdds);
+    runAnalysisPerCollisionWithUpc<true, false>(collisions, selectedLcCandidatesMl, bcs, ft0s, fv0as, fdds);
   }
   PROCESS_SWITCH(HfTaskUpcLc, processDataWithMlWithUpc, "Process real data with the ML method with UPC", false);
 
@@ -334,9 +389,37 @@ struct HfTaskUpcLc {
                              aod::FDDs const& fdds,
                              aod::Zdcs const& /*zdcs*/)
   {
-    runAnalysisPerCollisionDataWithUpc<false>(collisions, selectedLcCandidates, bcs, ft0s, fv0as, fdds);
+    runAnalysisPerCollisionWithUpc<false, false>(collisions, selectedLcCandidates, bcs, ft0s, fv0as, fdds);
   }
   PROCESS_SWITCH(HfTaskUpcLc, processDataStdWithUpc, "Process real data with the standard method with UPC", false);
+
+  void processMcWithMlWithUpc(soa::Join<aod::Collisions, aod::McCollisionLabels, aod::EvSels, aod::Mults> const& collisions,
+                              aod::BcFullInfos const& bcs,
+                              LcCandidatesMlMc const& selectedLcCandidatesMlMc,
+                              aod::McCollisions const&,
+                              aod::TracksWMc const&,
+                              aod::FT0s const& ft0s,
+                              aod::FV0As const& fv0as,
+                              aod::FDDs const& fdds,
+                              aod::Zdcs const& /*zdcs*/)
+  {
+    runAnalysisPerCollisionWithUpc<true, true>(collisions, selectedLcCandidatesMlMc, bcs, ft0s, fv0as, fdds);
+  }
+  PROCESS_SWITCH(HfTaskUpcLc, processMcWithMlWithUpc, "Process MC with the ML method with UPC", false);
+
+  void processMcStdWithUpc(soa::Join<aod::Collisions, aod::McCollisionLabels, aod::EvSels, aod::Mults> const& collisions,
+                           aod::BcFullInfos const& bcs,
+                           LcCandidatesMc const& selectedLcCandidatesMc,
+                           aod::McCollisions const&,
+                           aod::TracksWMc const&,
+                           aod::FT0s const& ft0s,
+                           aod::FV0As const& fv0as,
+                           aod::FDDs const& fdds,
+                           aod::Zdcs const& /*zdcs*/)
+  {
+    runAnalysisPerCollisionWithUpc<false, true>(collisions, selectedLcCandidatesMc, bcs, ft0s, fv0as, fdds);
+  }
+  PROCESS_SWITCH(HfTaskUpcLc, processMcStdWithUpc, "Process MC with the standard method with UPC", false);
 };
 
 WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)