Use invariant mass for rapidity calculations

Author: fchinu

PWGHF/D2H/Tasks/taskDs.cxx

@@ -596,27 +596,25 @@ struct HfTaskDs {
   {
 
     int id = o2::constants::physics::Pdg::kDS;
-    auto yCand = hfHelper.yDs(candidate);
     if (dataType == DataType::McDplusPrompt || dataType == DataType::McDplusNonPrompt || dataType == DataType::McDplusBkg) {
       id = o2::constants::physics::Pdg::kDPlus;
-      yCand = hfHelper.yDplus(candidate);
     } else if (dataType == DataType::McLcBkg) {
       id = o2::constants::physics::Pdg::kLambdaCPlus;
-      yCand = hfHelper.yLc(candidate);
     }
 
     auto indexMother = RecoDecay::getMother(mcParticles,
                                             candidate.template prong0_as<aod::TracksWMc>().template mcParticle_as<CandDsMcGen>(),
                                             id, true);
 
     if (indexMother != -1) {
-      if (yCandRecoMax >= 0. && std::abs(yCand) > yCandRecoMax) {
-        return;
-      }
 
       auto pt = candidate.pt(); // rec. level pT
 
       if (candidate.isSelDsToKKPi() >= selectionFlagDs) { // KKPi
+        auto yCand = candidate.y(hfHelper.invMassDsToKKPi(candidate));
+        if (yCandRecoMax >= 0. && std::abs(yCand) > yCandRecoMax) {
+          return;
+        }
         fillHisto(candidate, dataType);
         fillHistoKKPi<true, Coll>(candidate, dataType);
 
@@ -631,6 +629,10 @@ struct HfTaskDs {
         }
       }
       if (candidate.isSelDsToPiKK() >= selectionFlagDs) { // PiKK
+        auto yCand = candidate.y(hfHelper.invMassDsToPiKK(candidate));
+        if (yCandRecoMax >= 0. && std::abs(yCand) > yCandRecoMax) {
+          return;
+        }
         fillHisto(candidate, dataType);
         fillHistoPiKK<true, Coll>(candidate, dataType);
 
@@ -651,15 +653,17 @@ struct HfTaskDs {
   template <typename Coll, typename CandDs>
   void runDataAnalysisPerCandidate(CandDs const& candidate)
   {
-    if (yCandRecoMax >= 0. && std::abs(hfHelper.yDs(candidate)) > yCandRecoMax) {
-      return;
-    }
-
     if (candidate.isSelDsToKKPi() >= selectionFlagDs) { // KKPi
+      if (yCandRecoMax >= 0. && std::abs(candidate.y(hfHelper.invMassDsToKKPi(candidate))) > yCandRecoMax) {
+        return;
+      }
       fillHisto(candidate, DataType::Data);
       fillHistoKKPi<false, Coll>(candidate, DataType::Data);
     }
     if (candidate.isSelDsToPiKK() >= selectionFlagDs) { // PiKK
+      if (yCandRecoMax >= 0. && std::abs(candidate.y(hfHelper.invMassDsToPiKK(candidate))) > yCandRecoMax) {
+        return;
+      }
       fillHisto(candidate, DataType::Data);
       fillHistoPiKK<false, Coll>(candidate, DataType::Data);
     }
@@ -687,16 +691,19 @@ struct HfTaskDs {
       }
     }
     if (isBkg && fillMcBkgHistos) {
-      if (yCandRecoMax >= 0. && std::abs(hfHelper.yDs(candidate)) > yCandRecoMax) {
-        return;
-      }
 
       if (candidate.isSelDsToKKPi() >= selectionFlagDs || candidate.isSelDsToPiKK() >= selectionFlagDs) {
         if (candidate.isSelDsToKKPi() >= selectionFlagDs) { // KKPi
+          if (yCandRecoMax >= 0. && std::abs(candidate.y(hfHelper.invMassDsToKKPi(candidate))) > yCandRecoMax) {
+            return;
+          }
           fillHisto(candidate, DataType::McBkg);
           fillHistoKKPi<true, Coll>(candidate, DataType::McBkg);
         }
         if (candidate.isSelDsToPiKK() >= selectionFlagDs) { // PiKK
+          if (yCandRecoMax >= 0. && std::abs(candidate.y(hfHelper.invMassDsToPiKK(candidate))) > yCandRecoMax) {
+            return;
+          }
           fillHisto(candidate, DataType::McBkg);
           fillHistoPiKK<true, Coll>(candidate, DataType::McBkg);
         }

PWGHF/TableProducer/treeCreatorDsToKKPi.cxx

@@ -324,42 +324,41 @@ struct HfTreeCreatorDsToKKPi {
   template <bool doMc = false, int massHypo = 0, typename Coll, typename T>
   void fillCandidateTable(const T& candidate)
   {
-    int8_t flagMc = 0;
-    int8_t originMc = 0;
-    int8_t channelMc = 0;
-    int8_t isSwapped = massHypo; // 0 if KKPi, 1 if PiKK
-    float yCand = 0;
-    float eCand = 0;
-    float ctCand = 0;
+    float invMassDs = 0;
+    float deltaMassPhiKK = 0;
+    float absCos3PiKDs = 0;
+    if constexpr (massHypo == 0) {
+      invMassDs = hfHelper.invMassDsToKKPi(candidate);
+      deltaMassPhiKK = hfHelper.deltaMassPhiDsToKKPi(candidate);
+      absCos3PiKDs = hfHelper.absCos3PiKDsToKKPi(candidate);
+    } else if constexpr (massHypo == 1) {
+      invMassDs = hfHelper.invMassDsToPiKK(candidate);
+      deltaMassPhiKK = hfHelper.deltaMassPhiDsToPiKK(candidate);
+      absCos3PiKDs = hfHelper.absCos3PiKDsToPiKK(candidate);
+    }
+
+    int8_t flagMc{0};
+    int8_t originMc{0};
+    int8_t channelMc{0};
+    int8_t isSwapped{massHypo}; // 0 if KKPi, 1 if PiKK
+    float yCand = isSwapped ? candidate.y(invMassDs) : candidate.y(invMassDs);
+    float eCand{0.f};
+    float ctCand{0.f};
     if constexpr (doMc) {
       flagMc = candidate.flagMcMatchRec();
       originMc = candidate.originMcRec();
       channelMc = candidate.flagMcDecayChanRec();
       isSwapped = candidate.isCandidateSwapped();
+      yCand = isSwapped ? candidate.y(invMassDs) : candidate.y(invMassDs);
       if (fillDplusMc && candidate.flagMcDecayChanRec() == channelsResonant[Mother::Dplus][decayChannel]) {
-        yCand = hfHelper.yDplus(candidate);
         eCand = hfHelper.eDplus(candidate);
         ctCand = hfHelper.ctDplus(candidate);
       } else {
-        yCand = hfHelper.yDs(candidate);
         eCand = hfHelper.eDs(candidate);
         ctCand = hfHelper.ctDs(candidate);
       }
     }
 
-    float invMassDs = 0;
-    float deltaMassPhiKK = 0;
-    float absCos3PiKDs = 0;
-    if constexpr (massHypo == 0) {
-      invMassDs = hfHelper.invMassDsToKKPi(candidate);
-      deltaMassPhiKK = hfHelper.deltaMassPhiDsToKKPi(candidate);
-      absCos3PiKDs = hfHelper.absCos3PiKDsToKKPi(candidate);
-    } else if constexpr (massHypo == 1) {
-      invMassDs = hfHelper.invMassDsToPiKK(candidate);
-      deltaMassPhiKK = hfHelper.deltaMassPhiDsToPiKK(candidate);
-      absCos3PiKDs = hfHelper.absCos3PiKDsToPiKK(candidate);
-    }
-
     auto const& collision = candidate.template collision_as<Coll>();
     float centrality = o2::hf_centrality::getCentralityColl(collision);
 
@@ -617,7 +616,7 @@ struct HfTreeCreatorDsToKKPi {
         particle.pt(),
         particle.eta(),
         particle.phi(),
-        RecoDecay::y(particle.pVector(), o2::constants::physics::MassDS),
+        std::abs(particle.pdgCode()) == o2::constants::physics::Pdg::kDS ? RecoDecay::y(particle.pVector(), o2::constants::physics::MassDS) : RecoDecay::y(particle.pVector(), o2::constants::physics::MassDPlus),
         particle.flagMcMatchGen(),
         particle.originMcGen());
     }