From b8ad7d54faa132628820d8686131646ee0256f23 Mon Sep 17 00:00:00 2001
From: feisenhu <53603353+feisenhu@users.noreply.github.com>
Date: Sun, 8 Sep 2024 20:27:09 +0200
Subject: [PATCH] [PWGDQ] Add track propagation for pairs in dqefficiency task
 (#7614)

---
 PWGDQ/Tasks/dqEfficiency_withAssoc.cxx | 20 +++++++++++++++++---
 1 file changed, 17 insertions(+), 3 deletions(-)

diff --git a/PWGDQ/Tasks/dqEfficiency_withAssoc.cxx b/PWGDQ/Tasks/dqEfficiency_withAssoc.cxx
index b288c1f3974..7d6e0bbc926 100644
--- a/PWGDQ/Tasks/dqEfficiency_withAssoc.cxx
+++ b/PWGDQ/Tasks/dqEfficiency_withAssoc.cxx
@@ -874,6 +874,8 @@ struct AnalysisPrefilterSelection {
   Configurable<std::string> fConfigPrefilterTrackCut{"cfgPrefilterTrackCut", "", "Prefilter track cut"};
   Configurable<std::string> fConfigPrefilterPairCut{"cfgPrefilterPairCut", "", "Prefilter pair cut"};
   Configurable<std::string> fConfigTrackCuts{"cfgTrackCuts", "", "Track cuts for which to run the prefilter"};
+  // Track related options
+  Configurable<bool> fPropTrack{"cfgPropTrack", true, "Propgate tracks to associated collision to recalculate DCA and momentum vector"};
 
   std::map<uint32_t, uint32_t> fPrefilterMap;
   AnalysisCompositeCut* fPairCut;
@@ -938,8 +940,8 @@ struct AnalysisPrefilterSelection {
     VarManager::SetupTwoProngFwdDCAFitter(5.0f, true, 200.0f, 1.0e-3f, 0.9f, true);
   }
 
-  template <uint32_t TTrackFillMap, typename TTracks>
-  void runPrefilter(soa::Join<aod::ReducedTracksAssoc, aod::BarrelTrackCuts> const& assocs, TTracks const& /*tracks*/)
+  template <uint32_t TTrackFillMap, typename TEvent, typename TTracks>
+  void runPrefilter(TEvent const& event, soa::Join<aod::ReducedTracksAssoc, aod::BarrelTrackCuts> const& assocs, TTracks const& /*tracks*/)
   {
 
     for (auto& [assoc1, assoc2] : o2::soa::combinations(assocs, assocs)) {
@@ -964,6 +966,9 @@ struct AnalysisPrefilterSelection {
 
       // compute pair quantities
       VarManager::FillPair<VarManager::kDecayToEE, TTrackFillMap>(track1, track2);
+      if (fPropTrack) {
+        VarManager::FillPairCollision<VarManager::kDecayToEE, TTrackFillMap>(event, track1, track2);
+      }
       // if the pair fullfils the criteria, add an entry into the prefilter map for the two tracks
       if (fPairCut->IsSelected(VarManager::fgValues)) {
         if (fPrefilterMap.find(track1.globalIndex()) == fPrefilterMap.end() && track1Candidate > 0) {
@@ -984,7 +989,7 @@ struct AnalysisPrefilterSelection {
     for (auto& event : events) {
       auto groupedAssocs = assocs.sliceBy(trackAssocsPerCollision, event.globalIndex());
       if (groupedAssocs.size() > 1) {
-        runPrefilter<gkTrackFillMap>(groupedAssocs, tracks);
+        runPrefilter<gkTrackFillMap>(event, groupedAssocs, tracks);
       }
     }
     uint32_t mymap = -1;
@@ -1056,6 +1061,9 @@ struct AnalysisSameEventPairing {
   Configurable<std::string> fConfigMCGenSignals{"cfgBarrelMCGenSignals", "", "Comma separated list of MC signals (generated)"};
   Configurable<bool> fConfigSkimSignalOnly{"fConfigSkimSignalOnly", false, "Configurable to select only matched candidates"};
 
+  // Track related options
+  Configurable<bool> fPropTrack{"cfgPropTrack", true, "Propgate tracks to associated collision to recalculate DCA and momentum vector"};
+
   Service<o2::ccdb::BasicCCDBManager> fCCDB;
 
   // Filter filterEventSelected = aod::dqanalysisflags::isEventSelected & uint32_t(1);
@@ -1442,6 +1450,9 @@ struct AnalysisSameEventPairing {
           }
 
           VarManager::FillPair<TPairType, TTrackFillMap>(t1, t2);
+          if (fPropTrack) {
+            VarManager::FillPairCollision<TPairType, TTrackFillMap>(event, t1, t2);
+          }
           if constexpr (TTwoProngFitter) {
             VarManager::FillPairVertexing<TPairType, TEventFillMap, TTrackFillMap>(event, t1, t2, fConfigPropToPCA);
           }
@@ -1490,6 +1501,9 @@ struct AnalysisSameEventPairing {
           }
 
           VarManager::FillPair<TPairType, TTrackFillMap>(t1, t2);
+          if (fPropTrack) {
+            VarManager::FillPairCollision<TPairType, TTrackFillMap>(event, t1, t2);
+          }
           if constexpr (TTwoProngFitter) {
             VarManager::FillPairVertexing<TPairType, TEventFillMap, TTrackFillMap>(event, t1, t2, fConfigPropToPCA);
           }