Refactor and complete the charm femtoscopy framework

Author: zhangbiao-phy

PWGCF/DataModel/FemtoDerived.h

@@ -37,13 +37,6 @@ namespace femtoDream
 /// \tparam partOne Type of particle 1 (Track/V0/Cascade/...)
 /// \tparam partTwo Type of particle 2 (Track/V0/Cascade/...)
 
-enum ProngCharmHadron {
-  Prong0 = 0,
-  Prong1 = 1,
-  Prong2 = 2,
-  Nprongs = 3
-};
-
 template <o2::aod::femtodreamparticle::ParticleType partOne, o2::aod::femtodreamparticle::ParticleType partTwo>
 class FemtoDreamDetaDphiStar
 {
@@ -118,8 +111,14 @@ class FemtoDreamDetaDphiStar
         }
       }
     }
-    if constexpr (mPartOneType == o2::aod::femtodreamparticle::ParticleType::kTrack && mPartTwoType == o2::aod::femtodreamparticle::ParticleType::kCharmHadron) {
-      for (int i = 0; i < Nprongs; i++) {
+    if constexpr (mPartOneType == o2::aod::femtodreamparticle::ParticleType::kTrack && (mPartTwoType == o2::aod::femtodreamparticle::ParticleType::kCharmHadron3Prong || mPartTwoType == o2::aod::femtodreamparticle::ParticleType::kCharmHadron2Prong)) {
+      int nProng = 0;
+      if (mPartTwoType == o2::aod::femtodreamparticle::ParticleType::kCharmHadron3Prong)
+        nProng = 3;
+      if (mPartTwoType == o2::aod::femtodreamparticle::ParticleType::kCharmHadron2Prong)
+        nProng = 2;
+
+      for (int i = 0; i < nProng; i++) {
         std::string dirName = static_cast<std::string>(dirNames[2]);
         histdetadpi[i][0] = mHistogramRegistry->add<TH2>((dirName + static_cast<std::string>(histNames[0][i]) + static_cast<std::string>(histNameSEorME[meORse])).c_str(), "; #Delta #eta; #Delta #phi^{*}", kTH2F, {{100, -0.15, 0.15}, {100, -0.15, 0.15}});
         histdetadpi[i][1] = mHistogramRegistry->add<TH2>((dirName + static_cast<std::string>(histNames[1][i]) + static_cast<std::string>(histNameSEorME[meORse])).c_str(), "; #Delta #eta; #Delta #phi^{*}", kTH2F, {{100, -0.15, 0.15}, {100, -0.15, 0.15}});
@@ -495,50 +494,102 @@ class FemtoDreamDetaDphiStar
       }
 
       return pass;
-    } else if constexpr (mPartOneType == o2::aod::femtodreamparticle::ParticleType::kTrack && mPartTwoType == o2::aod::femtodreamparticle::ParticleType::kCharmHadron) {
+    } else if constexpr (mPartOneType == o2::aod::femtodreamparticle::ParticleType::kTrack && (mPartTwoType == o2::aod::femtodreamparticle::ParticleType::kCharmHadron3Prong || mPartTwoType == o2::aod::femtodreamparticle::ParticleType::kCharmHadronDstar)) { // to be fixed for Dstar
       // check if provided particles are in agreement with the class instantiation
-      if (part2.candidateSelFlag() < o2::aod::fdhf::lcToPKPi) {
+      if (!part2.candidateSelFlag()) {
         LOG(fatal) << "FemtoDreamDetaDphiStar: passed arguments don't agree with FemtoDreamDetaDphiStar instantiation! Please provide Charm Hadron candidates.";
         return false;
       }
 
       bool pass = false;
 
-      for (int i = 0; i < Nprongs; ++i) {
+      for (int i = 0; i < 3; ++i) {
         double deta, dphiAvg, dphi_AT_PV, dphi_AT_SpecificRadii, daughterEta, daughterPhi;
         bool sameCharge = false;
-        daughterEta = -999.;
-        daughterPhi = -999.;
-
-        switch (i) {
-          case Prong0:
-            daughterEta = part2.prong0Eta();
-            daughterPhi = part2.prong0Phi();
-            deta = part1.eta() - daughterEta;
-            dphi_AT_PV = part1.phi() - daughterPhi;
-            dphi_AT_SpecificRadii = PhiAtSpecificRadiiTPC(part1, radiiTPC) - PhiAtSpecificRadiiTPC<true, 0>(part2, radiiTPC);
-            dphiAvg = AveragePhiStar<true>(part1, part2, 0, &sameCharge);
-            // histdetadpi[0][0]->Fill(deta, dphiAvg);
-            break;
-          case Prong1:
-            daughterEta = part2.prong1Eta();
-            daughterPhi = part2.prong1Phi();
-            deta = part1.eta() - daughterEta;
-            dphi_AT_PV = part1.phi() - daughterPhi;
-            dphi_AT_SpecificRadii = PhiAtSpecificRadiiTPC(part1, radiiTPC) - PhiAtSpecificRadiiTPC<true, 1>(part2, radiiTPC);
-            dphiAvg = AveragePhiStar<true>(part1, part2, 1, &sameCharge);
-            // histdetadpi[1][0]->Fill(deta, dphiAvg);
-            break;
-          case Prong2:
-            daughterEta = part2.prong2Eta();
-            daughterPhi = part2.prong2Phi();
-            deta = part1.eta() - daughterEta;
-            dphi_AT_PV = part1.phi() - daughterPhi;
-            dphi_AT_SpecificRadii = PhiAtSpecificRadiiTPC(part1, radiiTPC) - PhiAtSpecificRadiiTPC<true, 2>(part2, radiiTPC);
-            dphiAvg = AveragePhiStar<true>(part1, part2, 2, &sameCharge);
-            // histdetadpi[2][0]->Fill(deta, dphiAvg);
-            break;
+        daughterEta = part2.prong2Eta();
+        daughterPhi = part2.prong2Phi();
+        deta = part1.eta() - daughterEta;
+        dphi_AT_PV = part1.phi() - daughterPhi;
+        dphi_AT_SpecificRadii = PhiAtSpecificRadiiTPC(part1, radiiTPC) - PhiAtSpecificRadiiTPC<true, 2>(part2, radiiTPC);
+        dphiAvg = AveragePhiStar<true>(part1, part2, 2, &sameCharge);
+
+        if (Q3 == 999) {
+          histdetadpi[i][0]->Fill(deta, dphiAvg);
+          histdetadpi[i][2]->Fill(deta, dphi_AT_PV);
+          if (fillQA) {
+            histdetadpi_eta[i]->Fill(deta, dphiAvg, part1.eta(), daughterEta);
+            histdetadpi_phi[i]->Fill(deta, dphiAvg, part1.phi(), daughterPhi);
+          }
+        } else if (Q3 < upperQ3LimitForPlotting) {
+          histdetadpi[i][0]->Fill(deta, dphiAvg);
+          histdetadpi[i][2]->Fill(deta, dphi_AT_PV);
+          if (fillQA) {
+            histdetadpi_eta[i]->Fill(deta, dphiAvg, part1.eta(), daughterEta);
+            histdetadpi_phi[i]->Fill(deta, dphiAvg, part1.phi(), daughterPhi);
+          }
         }
+
+        if (atWhichRadiiToSelect == 1) {
+          if (std::pow(dphiAvg, 2) / std::pow(deltaPhiMax, 2) + std::pow(deta, 2) / std::pow(deltaEtaMax, 2) < 1.) {
+            pass = true;
+          } else {
+            if (Q3 == 999) {
+              histdetadpi[i][1]->Fill(deta, dphiAvg);
+              histdetadpi[i][3]->Fill(deta, dphi_AT_PV);
+            } else if (Q3 < upperQ3LimitForPlotting) {
+              histdetadpi[i][1]->Fill(deta, dphiAvg);
+              histdetadpi[i][3]->Fill(deta, dphi_AT_PV);
+            }
+          }
+        } else if (atWhichRadiiToSelect == 0) {
+          if (std::pow(dphi_AT_PV, 2) / std::pow(deltaPhiMax, 2) + std::pow(deta, 2) / std::pow(deltaEtaMax, 2) < 1.) {
+            pass = true;
+          } else {
+            if (Q3 == 999) {
+              histdetadpi[i][1]->Fill(deta, dphiAvg);
+              histdetadpi[i][3]->Fill(deta, dphi_AT_PV);
+            } else if (Q3 < upperQ3LimitForPlotting) {
+              histdetadpi[i][1]->Fill(deta, dphiAvg);
+              histdetadpi[i][3]->Fill(deta, dphi_AT_PV);
+            }
+          }
+        } else if (atWhichRadiiToSelect == 2) {
+          if (std::pow(dphi_AT_SpecificRadii, 2) / std::pow(deltaPhiMax, 2) + std::pow(deta, 2) / std::pow(deltaEtaMax, 2) < 1.) {
+            pass = true;
+          } else {
+            if (Q3 == 999) {
+              histdetadpi[i][1]->Fill(deta, dphiAvg);
+              histdetadpi[i][3]->Fill(deta, dphi_AT_PV);
+            } else if (Q3 < upperQ3LimitForPlotting) {
+              histdetadpi[i][1]->Fill(deta, dphiAvg);
+              histdetadpi[i][3]->Fill(deta, dphi_AT_PV);
+            }
+          }
+        }
+      }
+
+      return pass;
+
+    } else if constexpr (mPartOneType == o2::aod::femtodreamparticle::ParticleType::kTrack && mPartTwoType == o2::aod::femtodreamparticle::ParticleType::kCharmHadron2Prong) {
+      // check if provided particles are in agreement with the class instantiation
+      if (!part2.candidateSelFlag()) {
+        LOG(fatal) << "FemtoDreamDetaDphiStar: passed arguments don't agree with FemtoDreamDetaDphiStar instantiation! Please provide Charm Hadron candidates.";
+        return false;
+      }
+
+      bool pass = false;
+
+      for (int i = 0; i < 2; ++i) {
+        double deta, dphiAvg, dphi_AT_PV, dphi_AT_SpecificRadii, daughterEta, daughterPhi;
+        bool sameCharge = false;
+        daughterEta = part2.prong1Eta();
+        daughterPhi = part2.prong1Phi();
+        deta = part1.eta() - daughterEta;
+        dphi_AT_PV = part1.phi() - daughterPhi;
+        dphi_AT_SpecificRadii = PhiAtSpecificRadiiTPC(part1, radiiTPC) - PhiAtSpecificRadiiTPC<true, 1>(part2, radiiTPC);
+        dphiAvg = AveragePhiStar<true>(part1, part2, 1, &sameCharge);
+        // histdetadpi[1][0]->Fill(deta, dphiAvg);
+
         if (Q3 == 999) {
           histdetadpi[i][0]->Fill(deta, dphiAvg);
           histdetadpi[i][2]->Fill(deta, dphi_AT_PV);
@@ -775,27 +826,25 @@ class FemtoDreamDetaDphiStar
   float PhiAtSpecificRadiiTPC(const T& part, float radii)
   {
     int charge = 0;
-    float phi0, pt;
+    float phi0 = 0.f;
+    float pt = 0.f;
+
     if constexpr (isHF) {
-      switch (prong) {
-        case Prong0:
-          charge = part.charge(); // charge calculation according to 3-prong decay, Lc^+ --> P^+ + K^- + pi^+
-          phi0 = part.prong0Phi();
-          pt = part.prong0Pt();
-          break;
-        case Prong1:
-          charge = -part.charge();
-          phi0 = part.prong1Phi();
-          pt = part.prong1Pt();
-          break;
-        case Prong2:
-          charge = part.charge();
-          phi0 = part.prong2Phi();
-          pt = part.prong2Pt();
-          break;
-        default:
-          // Handle invalid prong value if necessary
-          break;
+      static_assert(prong == 0 || prong == 1 || prong == 2,
+                    "PhiAtSpecificRadiiTPC: invalid prong index for HF");
+
+      if constexpr (prong == 0) {
+        charge = part.charge();
+        phi0 = part.prong0Phi();
+        pt = part.prong0Pt();
+      } else if constexpr (prong == 1) {
+        charge = -part.charge();
+        phi0 = part.prong1Phi();
+        pt = part.prong1Pt();
+      } else if constexpr (prong == 2) {
+        charge = part.charge();
+        phi0 = part.prong2Phi();
+        pt = part.prong2Pt();
       }
     } else {
       phi0 = part.phi();
@@ -833,42 +882,42 @@ class FemtoDreamDetaDphiStar
   int PhiAtRadiiTPCForHF(const T& part, std::vector<float>& tmpVec, int prong)
   {
     int charge = 0;
-    float pt = -999.;
-    float phi0 = -999.;
-    switch (prong) {
-      case Prong0:
-        pt = part.prong0Pt();
-        phi0 = part.prong0Phi();
+    if constexpr (mPartTwoType == o2::aod::femtodreamparticle::kCharmHadron3Prong) {
+      if (prong == 0 || prong == 2) {
         charge = part.charge();
-        break;
-      case Prong1:
-        pt = part.prong1Pt();
-        phi0 = part.prong1Phi();
+      } else if (prong == 1) {
         charge = -part.charge();
-        break;
-      case Prong2:
-        pt = part.prong2Pt();
-        phi0 = part.prong2Phi();
+      } else {
+        return 0;
+      }
+      for (size_t i = 0; i < 9; ++i) {
+        if (prong == 0) {
+          tmpVec.push_back(PhiAtSpecificRadiiTPC<true, 0>(part, tmpRadiiTPC[i]));
+        } else if (prong == 1) {
+          tmpVec.push_back(PhiAtSpecificRadiiTPC<true, 1>(part, tmpRadiiTPC[i]));
+        } else { // prong == 2
+          tmpVec.push_back(PhiAtSpecificRadiiTPC<true, 2>(part, tmpRadiiTPC[i]));
+        }
+      }
+
+    } else {
+      if (prong == 0) {
         charge = part.charge();
-        break;
-      default:
-        // Handle invalid prong value
-        break;
-    }
-    for (size_t i = 0; i < 9; i++) {
-      if (runOldVersion) {
-        tmpVec.push_back(phi0 - std::asin(0.3 * charge * 0.1 * magfield * tmpRadiiTPC[i] * 0.01 / (2. * pt)));
+      } else if (prong == 1) {
+        charge = -part.charge();
+      } else {
+        return 0;
       }
-      if (!runOldVersion) {
-        auto arg = 0.3 * charge * magfield * tmpRadiiTPC[i] * 0.01 / (2. * pt);
-        // for very low pT particles, this value goes outside of range -1 to 1 at at large tpc radius; asin fails
-        if (std::fabs(arg) < 1) {
-          tmpVec.push_back(phi0 - std::asin(0.3 * charge * magfield * tmpRadiiTPC[i] * 0.01 / (2. * pt)));
-        } else {
-          tmpVec.push_back(999);
+
+      for (size_t i = 0; i < 9; ++i) {
+        if (prong == 0) {
+          tmpVec.push_back(PhiAtSpecificRadiiTPC<true, 0>(part, tmpRadiiTPC[i]));
+        } else { // prong == 1
+          tmpVec.push_back(PhiAtSpecificRadiiTPC<true, 1>(part, tmpRadiiTPC[i]));
         }
       }
     }
+
     return charge;
   }
 

PWGCF/FemtoDream/Core/femtoDreamDetaDphiStar.h

@@ -17,6 +17,7 @@
 #define PWGCF_FEMTODREAM_CORE_FEMTODREAMPAIRCLEANER_H_
 
 #include "PWGCF/DataModel/FemtoDerived.h"
+
 #include "Framework/HistogramRegistry.h"
 
 using namespace o2::framework;
@@ -74,9 +75,9 @@ class FemtoDreamPairCleaner
         return true;
       }
       return false;
-    } else if constexpr (mPartOneType == o2::aod::femtodreamparticle::ParticleType::kTrack && mPartTwoType == o2::aod::femtodreamparticle::ParticleType::kCharmHadron) {
+    } else if constexpr (mPartOneType == o2::aod::femtodreamparticle::ParticleType::kTrack && (mPartTwoType == o2::aod::femtodreamparticle::ParticleType::kCharmHadron3Prong || mPartTwoType == o2::aod::femtodreamparticle::ParticleType::kCharmHadronDstar)) {
       /// Track-CharmHadron combination
-      if (part2.candidateSelFlag() < o2::aod::fdhf::lcToPKPi) {
+      if (!part2.candidateSelFlag()) {
         LOG(fatal) << "FemtoDreamPairCleaner: passed arguments don't agree with FemtoDreamPairCleaner instantiation! Please provide second argument Charm candidate.";
         return false;
       }
@@ -85,6 +86,17 @@ class FemtoDreamPairCleaner
         return true;
       }
       return false;
+    } else if constexpr (mPartOneType == o2::aod::femtodreamparticle::ParticleType::kTrack && mPartTwoType == o2::aod::femtodreamparticle::ParticleType::kCharmHadron2Prong) {
+      /// Track-CharmHadron combination
+      if (!part2.candidateSelFlag()) {
+        LOG(fatal) << "FemtoDreamPairCleaner: passed arguments don't agree with FemtoDreamPairCleaner instantiation! Please provide second argument Charm candidate.";
+        return false;
+      }
+
+      if (part1.trackId() != part2.prong0Id() && part1.trackId() != part2.prong1Id()) {
+        return true;
+      }
+      return false;
     } else if constexpr (mPartOneType == o2::aod::femtodreamparticle::ParticleType::kTrack && mPartTwoType == o2::aod::femtodreamparticle::ParticleType::kCascade) {
       /// Track-Cascade combination
       if (part2.partType() != o2::aod::femtodreamparticle::ParticleType::kCascade) {