[PWGCF] adding 2prong-2prong correlations (#8686)

Author: jikim1290

PWGCF/Tasks/correlations.cxx

@@ -9,6 +9,8 @@
 // granted to it by virtue of its status as an Intergovernmental Organization
 // or submit itself to any jurisdiction.
 #include <experimental/type_traits>
+#include <vector>
+#include <string>
 
 #include <TH1F.h>
 #include <cmath>
@@ -104,7 +106,7 @@ struct CorrelationTask {
   // This filter is applied to AOD and derived data (column names are identical)
   Filter collisionZVtxFilter = nabs(aod::collision::posZ) < cfgCutVertex;
   // This filter is only applied to AOD
-  Filter collisionVertexTypeFilter = (aod::collision::flags & (uint16_t)aod::collision::CollisionFlagsRun2::Run2VertexerTracks) == (uint16_t)aod::collision::CollisionFlagsRun2::Run2VertexerTracks;
+  Filter collisionVertexTypeFilter = (aod::collision::flags & static_cast<uint16_t>(aod::collision::CollisionFlagsRun2::Run2VertexerTracks)) == static_cast<uint16_t>(aod::collision::CollisionFlagsRun2::Run2VertexerTracks);
 
   // Track filters
   Filter trackFilter = (nabs(aod::track::eta) < cfgCutEta) && (aod::track::pt > cfgCutPt) && ((requireGlobalTrackInFilter()) || (aod::track::isGlobalTrackSDD == (uint8_t) true));
@@ -185,6 +187,9 @@ struct CorrelationTask {
     std::vector<AxisSpec> userAxis;
     if (cfgMassAxis != 0)
       userAxis.emplace_back(axisInvMass, "m (GeV/c^2)");
+    if (doprocessSame2Prong2Prong)
+      userAxis.emplace_back(axisInvMass, "m (GeV/c^2)");
+
     same.setObject(new CorrelationContainer("sameEvent", "sameEvent", corrAxis, effAxis, userAxis));
     mixed.setObject(new CorrelationContainer("mixedEvent", "mixedEvent", corrAxis, effAxis, userAxis));
 
@@ -241,7 +246,7 @@ struct CorrelationTask {
     for (auto& track1 : tracks1) {
       if constexpr (std::experimental::is_detected<hasInvMass, typename TTracks1::iterator>::value) {
         if constexpr (std::experimental::is_detected<hasDecay, typename TTracks1::iterator>::value) {
-          if (cfgDecayParticleMask != 0 && (cfgDecayParticleMask & (1u << (uint32_t)track1.decay())) == 0u)
+          if (cfgDecayParticleMask != 0 && (cfgDecayParticleMask & (1u << static_cast<uint32_t>(track1.decay()))) == 0u)
             continue;
         }
         registry.fill(HIST("invMass"), track1.invMass(), track1.pt(), multiplicity);
@@ -317,7 +322,7 @@ struct CorrelationTask {
       }
 
       if constexpr (std::experimental::is_detected<hasDecay, typename TTracks1::iterator>::value) {
-        if (cfgDecayParticleMask != 0 && (cfgDecayParticleMask & (1u << (uint32_t)track1.decay())) == 0u)
+        if (cfgDecayParticleMask != 0 && (cfgDecayParticleMask & (1u << static_cast<uint32_t>(track1.decay()))) == 0u)
           continue;
       }
 
@@ -375,28 +380,68 @@ struct CorrelationTask {
           }
         }
 
+        if constexpr (std::experimental::is_detected<hasDecay, typename TTracks2::iterator>::value) {
+          if (cfgDecayParticleMask != 0 && (cfgDecayParticleMask & (1u << static_cast<uint32_t>(track2.decay()))) == 0u)
+            continue;
+        }
+
+        if constexpr (std::experimental::is_detected<hasDecay, typename TTracks1::iterator>::value && std::experimental::is_detected<hasDecay, typename TTracks2::iterator>::value) {
+          if (doprocessSame2Prong2Prong && (track1.decay() == track2.decay() || track1.decay() > 1 || track2.decay() > 1)) {
+            continue;
+          }
+        } // D0 and anti-D0 selection
+
+        if constexpr (std::experimental::is_detected<hasProng0Id, typename TTracks1::iterator>::value) {
+          if constexpr (std::experimental::is_detected<hasProng0Id, typename TTracks2::iterator>::value) {
+            if (track1.cfTrackProng0Id() == track2.cfTrackProng0Id()) {
+              continue;
+            }
+          }
+          if constexpr (std::experimental::is_detected<hasProng1Id, typename TTracks2::iterator>::value) {
+            if (track1.cfTrackProng0Id() == track2.cfTrackProng1Id()) {
+              continue;
+            }
+          }
+        }
+
+        if constexpr (std::experimental::is_detected<hasProng1Id, typename TTracks1::iterator>::value) {
+          if constexpr (std::experimental::is_detected<hasProng0Id, typename TTracks2::iterator>::value) {
+            if (track1.cfTrackProng1Id() == track2.cfTrackProng0Id()) {
+              continue;
+            }
+          }
+          if constexpr (std::experimental::is_detected<hasProng1Id, typename TTracks2::iterator>::value) {
+            if (track1.cfTrackProng1Id() == track2.cfTrackProng1Id()) {
+              continue;
+            }
+          }
+        } // no shared prong for two mothers
+
         if (cfgPtOrder != 0 && track2.pt() >= track1.pt()) {
           continue;
         }
 
-        if (cfgAssociatedCharge != 0 && cfgAssociatedCharge * track2.sign() < 0) {
-          continue;
+        if constexpr (std::experimental::is_detected<hasSign, typename TTracks2::iterator>::value) {
+          if (cfgAssociatedCharge != 0 && cfgAssociatedCharge * track2.sign() < 0) {
+            continue;
+          }
         }
 
-        if constexpr (std::experimental::is_detected<hasSign, typename TTracks1::iterator>::value) {
+        if constexpr (std::experimental::is_detected<hasSign, typename TTracks1::iterator>::value && std::experimental::is_detected<hasSign, typename TTracks2::iterator>::value) {
           if (cfgPairCharge != 0 && cfgPairCharge * track1.sign() * track2.sign() < 0) {
             continue;
           }
         }
 
         if constexpr (std::is_same<TTracks1, TTracks2>::value) {
           if constexpr (step >= CorrelationContainer::kCFStepReconstructed) {
-            if (cfg.mPairCuts && mPairCuts.conversionCuts(track1, track2)) {
-              continue;
-            }
-
-            if (cfgTwoTrackCut > 0 && mPairCuts.twoTrackCut(track1, track2, magField)) {
-              continue;
+            if constexpr (std::experimental::is_detected<hasSign, typename TTracks1::iterator>::value && std::experimental::is_detected<hasSign, typename TTracks2::iterator>::value) {
+              if (cfg.mPairCuts && mPairCuts.conversionCuts(track1, track2)) {
+                continue;
+              }
+              if (cfgTwoTrackCut > 0 && mPairCuts.twoTrackCut(track1, track2, magField)) {
+                continue;
+              }
             }
           }
         }
@@ -417,7 +462,12 @@ struct CorrelationTask {
         }
 
         // last param is the weight
-        if (cfgMassAxis) {
+        if (cfgMassAxis && doprocessSame2Prong2Prong) {
+          if constexpr (std::experimental::is_detected<hasInvMass, typename TTracks1::iterator>::value && std::experimental::is_detected<hasInvMass, typename TTracks2::iterator>::value)
+            target->getPairHist()->Fill(step, track1.eta() - track2.eta(), track2.pt(), track1.pt(), multiplicity, deltaPhi, posZ, track2.invMass(), track1.invMass(), associatedWeight);
+          else
+            LOGF(fatal, "Can not fill mass axis without invMass column. \n no mass for two particles");
+        } else if (cfgMassAxis) {
           if constexpr (std::experimental::is_detected<hasInvMass, typename TTracks1::iterator>::value)
             target->getPairHist()->Fill(step, track1.eta() - track2.eta(), track2.pt(), track1.pt(), multiplicity, deltaPhi, posZ, track1.invMass(), associatedWeight);
           else
@@ -547,6 +597,30 @@ struct CorrelationTask {
   }
   PROCESS_SWITCH(CorrelationTask, processSame2ProngDerived, "Process same event on derived data", false);
 
+  void processSame2Prong2Prong(derivedCollisions::iterator const& collision, soa::Filtered<aod::CF2ProngTracks> const& p2tracks)
+  {
+    BinningTypeDerived configurableBinningDerived{{axisVertex, axisMultiplicity}, true}; // true is for 'ignore overflows' (true by default). Underflows and overflows will have bin -1.
+    if (cfgVerbosity > 0) {
+      LOGF(info, "processSame2ProngDerived: 2-prong candidates: %d | Vertex: %.1f | Multiplicity/Centrality: %.1f", p2tracks.size(), collision.posZ(), collision.multiplicity());
+    }
+    loadEfficiency(collision.timestamp());
+
+    const auto multiplicity = collision.multiplicity();
+
+    int bin = configurableBinningDerived.getBin({collision.posZ(), collision.multiplicity()});
+    registry.fill(HIST("eventcount_same"), bin);
+    fillQA(collision, multiplicity, p2tracks, p2tracks);
+
+    same->fillEvent(multiplicity, CorrelationContainer::kCFStepReconstructed);
+    fillCorrelations<CorrelationContainer::kCFStepReconstructed>(same, p2tracks, p2tracks, multiplicity, collision.posZ(), 0, 1.0f);
+
+    if (cfg.mEfficiencyAssociated || cfg.mEfficiencyTrigger) {
+      same->fillEvent(multiplicity, CorrelationContainer::kCFStepCorrected);
+      fillCorrelations<CorrelationContainer::kCFStepCorrected>(same, p2tracks, p2tracks, multiplicity, collision.posZ(), 0, 1.0f);
+    }
+  }
+  PROCESS_SWITCH(CorrelationTask, processSame2Prong2Prong, "Process same event on derived data", false);
+
   using BinningTypeAOD = ColumnBinningPolicy<aod::collision::PosZ, aod::cent::CentRun2V0M>;
   void processMixedAOD(aodCollisions& collisions, aodTracks const& tracks, aod::BCsWithTimestamps const&)
   {