Update treeCreatorOmegac0ToOmegaPi.cxx

Author: FDUEnrich

PWGHF/TableProducer/treeCreatorOmegac0ToOmegaPi.cxx

@@ -16,6 +16,7 @@
 /// \author Federica Zanone <federica.zanone@cern.ch>, Heidelberg University
 /// \author Yunfan Liu <yunfan.liu@cern.ch>, China University of Geosciences
 /// \author Fabio Catalano <fabio.catalano@cern.ch>, University of Houston
+/// \author Ruiqi Yin <ruiqi.yin@cern.ch>, Fudan University
 
 #include "Framework/AnalysisTask.h"
 #include "Framework/runDataProcessing.h"
@@ -151,6 +152,7 @@ DECLARE_SOA_COLUMN(OmegacChi2OverNdf, omegacChi2OverNdf, float);
 DECLARE_SOA_COLUMN(MassV0Chi2OverNdf, massV0Chi2OverNdf, float);
 DECLARE_SOA_COLUMN(MassCascChi2OverNdf, massCascChi2OverNdf, float);
 DECLARE_SOA_COLUMN(CascRejectInvmass, cascRejectInvmass, float);
+DECLARE_SOA_COLUMN(OutputMlOmegac, outputMlOmegac, float);
 } // namespace full
 
 DECLARE_SOA_TABLE(HfToOmegaPiEvs, "AOD", "HFTOOMEPIEV",
@@ -198,7 +200,7 @@ DECLARE_SOA_TABLE(HfKfOmegacFulls, "AOD", "HFKFOMEGACFULL",
                   full::V0Ndf, full::CascNdf, full::OmegacNdf,
                   full::MassV0Ndf, full::MassCascNdf,
                   full::V0Chi2OverNdf, full::CascChi2OverNdf, full::OmegacChi2OverNdf,
-                  full::MassV0Chi2OverNdf, full::MassCascChi2OverNdf, full::CascRejectInvmass,
+                  full::MassV0Chi2OverNdf, full::MassCascChi2OverNdf, full::CascRejectInvmass, full::OutputMlOmegac,
                   full::FlagMcMatchRec, full::OriginRec, full::CollisionMatched, hf_track_index::HFflag);
 
 DECLARE_SOA_TABLE(HfKfOmegacLites, "AOD", "HFKFOMEGACLITE",
@@ -214,14 +216,18 @@ DECLARE_SOA_TABLE(HfKfOmegacLites, "AOD", "HFKFOMEGACLITE",
                   full::V0Chi2OverNdf, full::CascChi2OverNdf, full::OmegacChi2OverNdf,
                   full::CascRejectInvmass,
                   full::FlagMcMatchRec, full::OriginRec, full::CollisionMatched, hf_track_index::HFflag);
+
+DECLARE_SOA_TABLE(HfKfPbPbOmegacMl, "AOD", "HFKFPBPBOMEGACML",
+                  full::InvMassCharmBaryon, full::KfptOmegac, full::KfptPiFromOmegac, full::OutputMlOmegac, full::FlagMcMatchRec, full::OriginRec, full::CollisionMatched, hf_track_index::HFflag);
 } // namespace o2::aod
 
 /// Writes the full information in an output TTree
-struct HfTreeCreatorOmegacToOmegaPi { // FIXME: Rename to HfTreeCreatorOmegac0ToOmegaPi
+struct HfTreeCreatorOmegac0ToOmegaPi {
 
   Produces<o2::aod::HfOmegac0ToOmegaPiLites> rowCandidateLite;
   Produces<o2::aod::HfKfOmegacFulls> rowKfCandidateFull;
   Produces<o2::aod::HfKfOmegacLites> rowKfCandidateLite;
+  Produces<o2::aod::HfKfPbPbOmegacMl> rowKfPbPbOmegacMl;
   Produces<o2::aod::HfToOmegaPiEvs> rowEv;
 
   Configurable<float> zPvCut{"zPvCut", 10., "Cut on absolute value of primary vertex z coordinate"};
@@ -381,6 +387,7 @@ struct HfTreeCreatorOmegacToOmegaPi { // FIXME: Rename to HfTreeCreatorOmegac0To
         candidate.massV0Chi2OverNdf(),
         candidate.massCascChi2OverNdf(),
         candidate.cascRejectInvmass(),
+        candidate.mlValueOmegac(),
         flagMc,
         originMc,
         collisionMatched,
@@ -431,6 +438,23 @@ struct HfTreeCreatorOmegacToOmegaPi { // FIXME: Rename to HfTreeCreatorOmegac0To
     }
   } // fillKfCandidateLite end
 
+  template <typename T>
+  void fillKfPbPbCandidateMl(const T& candidate, int8_t flagMc, int8_t originMc, bool collisionMatched)
+  {
+    if (candidate.resultSelections() && candidate.statusPidCharmBaryon() && candidate.statusInvMassLambda() && candidate.statusInvMassCascade() && candidate.statusInvMassCharmBaryon()) {
+
+      rowKfPbPbOmegacMl(
+        candidate.invMassCharmBaryon(),
+        candidate.kfptOmegac(),
+        candidate.kfptPiFromOmegac(),
+        candidate.mlValueOmegac(),
+        flagMc,
+        originMc,
+        collisionMatched,
+        candidate.hfflag());
+    }
+  } // fillPbPbCandidateMl end
+
   void processDataLite(Colls const& collisions, Tracks const&,
                        soa::Filtered<soa::Join<aod::HfCandToOmegaPi, aod::HfSelToOmegaPi>> const& candidates)
   {
@@ -446,7 +470,7 @@ struct HfTreeCreatorOmegacToOmegaPi { // FIXME: Rename to HfTreeCreatorOmegac0To
       fillCandidateLite(candidate, -7, RecoDecay::OriginType::None, false);
     }
   }
-  PROCESS_SWITCH(HfTreeCreatorOmegacToOmegaPi, processDataLite, "Process data", true);
+  PROCESS_SWITCH(HfTreeCreatorOmegac0ToOmegaPi, processDataLite, "Process data", false);
 
   void processKfDataFull(Colls const& collisions, Tracks const&, CandKfSel const& candidates)
   {
@@ -462,7 +486,7 @@ struct HfTreeCreatorOmegacToOmegaPi { // FIXME: Rename to HfTreeCreatorOmegac0To
       fillKfCandidate(candidate, -7, RecoDecay::OriginType::None, false);
     }
   }
-  PROCESS_SWITCH(HfTreeCreatorOmegacToOmegaPi, processKfDataFull, "Process KF data", false);
+  PROCESS_SWITCH(HfTreeCreatorOmegac0ToOmegaPi, processKfDataFull, "Process KF data", false);
 
   void processKfDataLite(Colls const& collisions, Tracks const&, CandKfSel const& candidates)
   {
@@ -478,7 +502,23 @@ struct HfTreeCreatorOmegacToOmegaPi { // FIXME: Rename to HfTreeCreatorOmegac0To
       fillKfCandidateLite(candidate, -7, RecoDecay::OriginType::None, false);
     }
   }
-  PROCESS_SWITCH(HfTreeCreatorOmegacToOmegaPi, processKfDataLite, "Process KF data Lite", false);
+  PROCESS_SWITCH(HfTreeCreatorOmegac0ToOmegaPi, processKfDataLite, "Process KF data Lite", false);
+
+  void processKfPbPbCandidateMl(Colls const& collisions, Tracks const&, CandKfSel const& candidates)
+  {
+    // Filling event properties
+    rowEv.reserve(collisions.size());
+    for (const auto& collision : collisions) {
+      fillEvent(collision, zPvCut);
+    }
+
+    // Filling candidate properties
+    rowKfCandidateFull.reserve(candidates.size());
+    for (const auto& candidate : candidates) {
+      fillKfPbPbCandidateMl(candidate, -7, RecoDecay::OriginType::None, false);
+    }
+  }
+  PROCESS_SWITCH(HfTreeCreatorOmegac0ToOmegaPi, processKfPbPbCandidateMl, "Process KF data ML", true);
 
   void processMcLite(Colls const& collisions, Tracks const&,
                      soa::Filtered<soa::Join<aod::HfCandToOmegaPi, aod::HfSelToOmegaPi, aod::HfToOmegaPiMCRec>> const& candidates)
@@ -495,7 +535,7 @@ struct HfTreeCreatorOmegacToOmegaPi { // FIXME: Rename to HfTreeCreatorOmegac0To
       fillCandidateLite(candidate, candidate.flagMcMatchRec(), candidate.originRec(), candidate.collisionMatched());
     }
   }
-  PROCESS_SWITCH(HfTreeCreatorOmegacToOmegaPi, processMcLite, "Process MC", false);
+  PROCESS_SWITCH(HfTreeCreatorOmegac0ToOmegaPi, processMcLite, "Process MC", false);
 
   void processKFMcFull(Colls const& collisions, Tracks const&, CascKfMcSel const& candidates)
   {
@@ -517,7 +557,7 @@ struct HfTreeCreatorOmegacToOmegaPi { // FIXME: Rename to HfTreeCreatorOmegac0To
       }
     }
   }
-  PROCESS_SWITCH(HfTreeCreatorOmegacToOmegaPi, processKFMcFull, "Process KF MC", false);
+  PROCESS_SWITCH(HfTreeCreatorOmegac0ToOmegaPi, processKFMcFull, "Process KF MC", false);
 
   void processKFMcLite(Colls const& collisions, Tracks const&, CascKfMcSel const& candidates)
   {
@@ -539,12 +579,12 @@ struct HfTreeCreatorOmegacToOmegaPi { // FIXME: Rename to HfTreeCreatorOmegac0To
       }
     }
   }
-  PROCESS_SWITCH(HfTreeCreatorOmegacToOmegaPi, processKFMcLite, "Process KF MC Lite", false);
+  PROCESS_SWITCH(HfTreeCreatorOmegac0ToOmegaPi, processKFMcLite, "Process KF MC Lite", false);
 
 }; // end of struct
 
 WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)
 {
   return WorkflowSpec{
-    adaptAnalysisTask<HfTreeCreatorOmegacToOmegaPi>(cfgc)};
+    adaptAnalysisTask<HfTreeCreatorOmegac0ToOmegaPi>(cfgc)};
 }