PWGHF: optimize the code to add Kf to D0

Author: lupengzhong

PWGHF/D2H/Tasks/taskD0.cxx

@@ -28,9 +28,6 @@ using namespace o2::framework::expressions;
 using namespace o2::aod::hf_cand_2prong;
 using namespace o2::analysis::hf_cuts_d0_to_pi_k;
 
-constexpr static int useDCAFitterN = 0;
-constexpr static int useKFParticle = 1;
-
 /// D0 analysis task
 struct HfTaskD0 {
   Configurable<int> selectionFlagD0{"selectionFlagD0", 1, "Selection Flag for D0"};
@@ -166,8 +163,8 @@ struct HfTaskD0 {
     registry.add("hDecLengthxyVsPtSig", "2-prong candidates;decay length xy (cm) vs #it{p}_{T} for signal;entries", {HistType::kTH2F, {{800, 0., 4.}, {vbins, "#it{p}_{T} (GeV/#it{c})"}}});
   }
 
-  template <int ReconstructionType, typename THfCand2Prong, typename THfCand2ProngSel>
-  void processData(THfCand2Prong& candidates, THfCand2ProngSel& selectedD0CandidatesSets)
+  template <int ReconstructionType, typename THfCand2ProngSel>
+  void processData(THfCand2ProngSel& selectedD0CandidatesSets)
   {
     for (auto& candidate : selectedD0CandidatesSets) {
       if (!(candidate.hfflag() & 1 << DecayType::D0ToPiK)) {
@@ -178,14 +175,13 @@ struct HfTaskD0 {
       }
 
       float massD0, massD0bar;
-      if constexpr (ReconstructionType == useDCAFitterN) {
+      if constexpr (ReconstructionType == o2::aod::hf_cand::useKFParticle) {
+        massD0 = candidate.kfGeoMassD0();
+        massD0bar = candidate.kfGeoMassD0bar();
+      } else {
         massD0 = invMassD0ToPiK(candidate);
         massD0bar = invMassD0barToKPi(candidate);
       }
-      if constexpr (ReconstructionType == useKFParticle) {
-        massD0 = candidate.kfGeoMass_DZero();
-        massD0bar = candidate.kfGeoMass_DZeroBar();
-      }
       auto ptCandidate = candidate.pt();
 
       if (candidate.isSelD0() >= selectionFlagD0) {
@@ -228,20 +224,20 @@ struct HfTaskD0 {
       registry.fill(HIST("hCPAXYFinerBinning"), candidate.cpaXY(), ptCandidate);
     }
   }
-  void processWithDCAFitterN(soa::Join<aod::HfCand2Prong, aod::HfSelD0> const& candidates)
+  void processWithDCAFitterN(soa::Join<aod::HfCand2Prong, aod::HfSelD0> const&)
   {
-    processData<0>(candidates, selectedD0Candidates);
+    processData<0>(selectedD0Candidates);
   }
-  void processWithKFParticle(soa::Join<aod::HfCand2Prong, aod::HfCand2ProngKF, aod::HfSelD0> const& candidates)
+  PROCESS_SWITCH(HfTaskD0, processWithDCAFitterN, "process taskD0 with DCAFitterN", true);
+
+  void processWithKFParticle(soa::Join<aod::HfCand2Prong, aod::HfCand2ProngKF, aod::HfSelD0> const&)
   {
-    processData<1>(candidates, selectedD0CandidatesKF);
+    processData<1>(selectedD0CandidatesKF);
   }
-
-  PROCESS_SWITCH(HfTaskD0, processWithDCAFitterN, "process taskD0 with DCAFitterN", true);
   PROCESS_SWITCH(HfTaskD0, processWithKFParticle, "process taskD0 with KFParticle", false);
 
-  template <int ReconstructionType, typename THfCand2Prong, typename THfCand2ProngFlag>
-  void processMc(THfCand2Prong& candidates, THfCand2ProngFlag& recoFlag2ProngSets,
+  template <int ReconstructionType, typename THfCand2ProngFlag>
+  void processMc(THfCand2ProngFlag& recoFlag2ProngSets,
                  soa::Join<aod::McParticles, aod::HfCand2ProngMcGen> const& particlesMC, aod::TracksWMc const& tracks)
   {
     // MC rec.
@@ -254,14 +250,13 @@ struct HfTaskD0 {
         continue;
       }
       float massD0, massD0bar;
-      if constexpr (ReconstructionType == useDCAFitterN) {
+      if constexpr (ReconstructionType == o2::aod::hf_cand::useKFParticle) {
+        massD0 = candidate.kfGeoMassD0();
+        massD0bar = candidate.kfGeoMassD0bar();
+      } else {
         massD0 = invMassD0ToPiK(candidate);
         massD0bar = invMassD0barToKPi(candidate);
       }
-      if constexpr (ReconstructionType == useKFParticle) {
-        massD0 = candidate.kfGeoMass_DZero();
-        massD0bar = candidate.kfGeoMass_DZeroBar();
-      }
       if (std::abs(candidate.flagMcMatchRec()) == 1 << DecayType::D0ToPiK) {
         // Get the corresponding MC particle.
         auto indexMother = RecoDecay::getMother(particlesMC, candidate.template prong0_as<aod::TracksWMc>().template mcParticle_as<soa::Join<aod::McParticles, aod::HfCand2ProngMcGen>>(), pdg::Code::kD0, true);
@@ -435,19 +430,16 @@ struct HfTaskD0 {
     }
   }
 
-  void processMcWithDCAFitterN(soa::Join<aod::HfCand2Prong, aod::HfSelD0, aod::HfCand2ProngMcRec>& candidates,
-                               soa::Join<aod::McParticles, aod::HfCand2ProngMcGen> const& particlesMC, aod::TracksWMc const& tracks)
+  void processMcWithDCAFitterN(soa::Join<aod::McParticles, aod::HfCand2ProngMcGen> const& particlesMC, aod::TracksWMc const& tracks)
   {
-    processMc<0>(candidates, recoFlag2Prong, particlesMC, tracks);
+    processMc<0>(recoFlag2Prong, particlesMC, tracks);
   }
+  PROCESS_SWITCH(HfTaskD0, processMcWithDCAFitterN, "Process MC with DCAFitterN", false);
 
-  void processMcWithKFParticle(soa::Join<aod::HfCand2Prong, aod::HfCand2ProngKF, aod::HfSelD0, aod::HfCand2ProngMcRec>& candidates,
-                               soa::Join<aod::McParticles, aod::HfCand2ProngMcGen> const& particlesMC, aod::TracksWMc const& tracks)
+  void processMcWithKFParticle(soa::Join<aod::McParticles, aod::HfCand2ProngMcGen> const& particlesMC, aod::TracksWMc const& tracks)
   {
-    processMc<1>(candidates, recoFlag2ProngKF, particlesMC, tracks);
+    processMc<1>(recoFlag2ProngKF, particlesMC, tracks);
   }
-
-  PROCESS_SWITCH(HfTaskD0, processMcWithDCAFitterN, "Process MC with DCAFitterN", false);
   PROCESS_SWITCH(HfTaskD0, processMcWithKFParticle, "Process MC with KFParticle", false);
 };
 

PWGHF/DataModel/CandidateReconstructionTables.h

@@ -261,9 +261,6 @@ DECLARE_SOA_COLUMN(ZSecondaryVertex, zSecondaryVertex, float); //!
 DECLARE_SOA_DYNAMIC_COLUMN(RSecondaryVertex, rSecondaryVertex, //!
                            [](float xVtxS, float yVtxS) -> float { return RecoDecay::sqrtSumOfSquares(xVtxS, yVtxS); });
 DECLARE_SOA_COLUMN(Chi2PCA, chi2PCA, float); //! sum of (non-weighted) distances of the secondary vertex to its prongs
-DECLARE_SOA_COLUMN(KFTopChi2OverNDF, kfTopChi2OverNDF, float);     //!
-DECLARE_SOA_COLUMN(KFGeoMass_DZero, kfGeoMass_DZero, float);       //!
-DECLARE_SOA_COLUMN(KFGeoMass_DZeroBar, kfGeoMass_DZeroBar, float); //!
 // prong properties
 DECLARE_SOA_COLUMN(PxProng0, pxProng0, float); //!
 DECLARE_SOA_COLUMN(PyProng0, pyProng0, float); //!
@@ -343,6 +340,9 @@ DECLARE_SOA_DYNAMIC_COLUMN(Ct, ct, //!
                            [](float xVtxP, float yVtxP, float zVtxP, float xVtxS, float yVtxS, float zVtxS, float px, float py, float pz, double m) -> float { return RecoDecay::ct(array{px, py, pz}, RecoDecay::distance(array{xVtxP, yVtxP, zVtxP}, array{xVtxS, yVtxS, zVtxS}), m); });
 DECLARE_SOA_DYNAMIC_COLUMN(ImpactParameterXY, impactParameterXY, //!
                            [](float xVtxP, float yVtxP, float zVtxP, float xVtxS, float yVtxS, float zVtxS, float px, float py, float pz) -> float { return RecoDecay::impParXY(array{xVtxP, yVtxP, zVtxP}, array{xVtxS, yVtxS, zVtxS}, array{px, py, pz}); });
+
+constexpr static int useDCAFitterN = 0;
+constexpr static int useKFParticle = 1;
 } // namespace hf_cand
 
 // specific 2-prong decay properties
@@ -371,6 +371,10 @@ DECLARE_SOA_COLUMN(FlagMcMatchRec, flagMcMatchRec, int8_t); //! reconstruction l
 DECLARE_SOA_COLUMN(FlagMcMatchGen, flagMcMatchGen, int8_t); //! generator level
 DECLARE_SOA_COLUMN(OriginMcRec, originMcRec, int8_t);       //! particle origin, reconstruction level
 DECLARE_SOA_COLUMN(OriginMcGen, originMcGen, int8_t);       //! particle origin, generator level
+// KF related properties
+DECLARE_SOA_COLUMN(KfTopolChi2OverNdf, kfTopolChi2OverNdf, float); //! chi2overndf of the KFParticle topological constraint
+DECLARE_SOA_COLUMN(KfGeoMassD0, kfGeoMassD0, float);               //! mass of the D0 candidate from the KFParticle geometric fit
+DECLARE_SOA_COLUMN(KfGeoMassD0bar, kfGeoMassD0bar, float);         //! mass of the D0bar candidate from the KFParticle geometric fit
 
 // mapping of decay types
 enum DecayType { D0ToPiK = 0,
@@ -523,8 +527,8 @@ DECLARE_SOA_EXTENDED_TABLE_USER(HfCand2ProngExt, HfCand2ProngBase, "HFCAND2PEXT"
 using HfCand2Prong = HfCand2ProngExt;
 
 DECLARE_SOA_TABLE(HfCand2ProngKF, "AOD", "HFCAND2PKF",
-                  hf_cand::KFTopChi2OverNDF,
-                  hf_cand::KFGeoMass_DZero, hf_cand::KFGeoMass_DZeroBar);
+                  hf_cand_2prong::KfTopolChi2OverNdf,
+                  hf_cand_2prong::KfGeoMassD0, hf_cand_2prong::KfGeoMassD0bar);
 
 // table with results of reconstruction level MC matching
 DECLARE_SOA_TABLE(HfCand2ProngMcRec, "AOD", "HFCAND2PMCREC", //!

PWGHF/TableProducer/CMakeLists.txt

@@ -25,7 +25,7 @@ o2physics_add_dpl_workflow(refit-pv-dummy
 
 o2physics_add_dpl_workflow(candidate-creator-2prong
                     SOURCES candidateCreator2Prong.cxx
-                    PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2::DCAFitter  KFParticle::KFParticle
+                    PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2::DCAFitter KFParticle::KFParticle
                     COMPONENT_NAME Analysis)
 
 o2physics_add_dpl_workflow(candidate-creator-3prong