From 8e5476ab44ee6241b2a3b9c5942f665bda54a316 Mon Sep 17 00:00:00 2001
From: KaiCui <kaicui@mails.ccnu.edu.cn>
Date: Thu, 16 Jul 2026 13:57:26 +0200
Subject: [PATCH] Update hStrangeCorrelation.cxx

---
 .../Tasks/Strangeness/hStrangeCorrelation.cxx | 41 ++++++++++++-------
 1 file changed, 27 insertions(+), 14 deletions(-)

diff --git a/PWGLF/Tasks/Strangeness/hStrangeCorrelation.cxx b/PWGLF/Tasks/Strangeness/hStrangeCorrelation.cxx
index 177c3c29d01..4b6c530b523 100644
--- a/PWGLF/Tasks/Strangeness/hStrangeCorrelation.cxx
+++ b/PWGLF/Tasks/Strangeness/hStrangeCorrelation.cxx
@@ -128,6 +128,7 @@ struct HStrangeCorrelation {
     Configurable<bool> useCentralityinPrediction{"useCentralityinPrediction", false, "if true, use centrality instead of multiplisity"};
     Configurable<bool> doMirroringInDelataEta{"doMirroringInDelataEta", false, "if true, fill only positive delta eta and mirror the negative side in post processing, Adjust the delta axis!"};
     Configurable<bool> fillCorrelationHistWithMass{"fillCorrelationHistWithMass", false, "if true, fill correlation histograms with particle mass"};
+    Configurable<bool> doMassSpectrumCheck{"doMassSpectrumCheck", false, "if true, add and fill invariant-mass spectrum"};
   } masterConfigurations;
 
   // master analysis switches
@@ -380,6 +381,19 @@ struct HStrangeCorrelation {
   static constexpr int IndexLambda = 1;
   static constexpr int IndexAntiLambda = 2;
 
+  template <int Index, typename TV0>
+  static float getV0InvariantMass(TV0 const& v0)
+  {
+    static_assert(Index >= IndexK0 && Index <= IndexAntiLambda);
+    if constexpr (Index == IndexK0) {
+      return v0.mK0Short();
+    } else if constexpr (Index == IndexLambda) {
+      return v0.mLambda();
+    } else {
+      return v0.mAntiLambda();
+    }
+  }
+
   uint16_t doCorrelation;
   int mRunNumber;
   int mRunNumberZorro;
@@ -949,19 +963,7 @@ struct HStrangeCorrelation {
             }
             if (assocCandidate.compatible(Index, trackSelection.dEdxCompatibility) && (!masterConfigurations.doMCassociation || assocCandidate.mcTrue(Index)) && (!doAssocPhysicalPrimary || assocCandidate.mcPhysicalPrimary()) && !mixing && (masterConfigurations.fillCorrelationHistWithMass || (-massWindowConfigurations.maxPeakNSigma < assocCandidate.invMassNSigma(Index) && assocCandidate.invMassNSigma(Index) < +massWindowConfigurations.maxPeakNSigma))) {
               if (masterConfigurations.fillCorrelationHistWithMass) {
-                float massAssoc = 0.f;
-                switch (i) {
-                  case 0:
-                    massAssoc = assoc.mK0Short();
-                    break;
-                  case 1:
-                    massAssoc = assoc.mLambda();
-                    break;
-                  case 2:
-                    massAssoc = assoc.mAntiLambda();
-                    break;
-                }
-                binFillThn[1] = massAssoc;
+                binFillThn[1] = getV0InvariantMass<Index>(assoc);
               }
               fillCorrelationHistogram(histos.get<THn>(HIST("sameEvent/Signal/") + HIST(V0names[Index])), binFillThn, etaWeight, efficiency * efficiencyTrigg, totalEffUncert, purityTrigg, purityTriggErr);
               if (std::abs(deltaphi) < checks.towardDeltaEtaRange && doITSClustersQA) {
@@ -1907,7 +1909,15 @@ struct HStrangeCorrelation {
       histos.addClone("sameEvent/Signal/", "sameEvent/LeftBg/");
       histos.addClone("sameEvent/Signal/", "sameEvent/RightBg/");
     }
-
+    if (TESTBIT(doCorrelation, 0) && doprocessSameEventHV0s && masterConfigurations.doMassSpectrumCheck) {
+      histos.add("hK0ShortPtVsMass", "", kTH2F, {axesConfigurations.axisPtQA, axesConfigurations.axisK0ShortMass});
+    }
+    if (TESTBIT(doCorrelation, 1) && doprocessSameEventHV0s && masterConfigurations.doMassSpectrumCheck) {
+      histos.add("hLambdaPtVsMass", "", kTH2F, {axesConfigurations.axisPtQA, axesConfigurations.axisLambdaMass});
+    }
+    if (TESTBIT(doCorrelation, 2) && doprocessSameEventHV0s && masterConfigurations.doMassSpectrumCheck) {
+      histos.add("hAntiLambdaPtVsMass", "", kTH2F, {axesConfigurations.axisPtQA, axesConfigurations.axisLambdaMass});
+    }
     LOGF(info, "Init THnFs done");
     // mixed-event correlation functions
     if ((doprocessMixedEventHV0sInBuffer || doprocessMixedEventHCascadesInBuffer || doprocessMixedEventHV0s || doprocessMixedEventHCascades || doprocessMixedEventHPions || doprocessMixedEventHHadrons) && masterConfigurations.doFullCorrelationStudy) {
@@ -2369,6 +2379,9 @@ struct HStrangeCorrelation {
         if (v0.compatible(Index, trackSelection.dEdxCompatibility) && (!masterConfigurations.doMCassociation || v0.mcTrue(Index)) && (!doAssocPhysicalPrimary || v0.mcPhysicalPrimary()) && (!efficiencyFlags.applyEfficiencyCorrection || efficiency != 0)) {
           if ((TESTBIT(doCorrelation, Index)) && (masterConfigurations.doPPAnalysis || (TESTBIT(selMap, Index) && TESTBIT(selMap, Index + 3)))) {
             histos.fill(HIST("h3d") + HIST(V0names[Index]) + HIST("Spectrum"), v0Data.pt(), cent, v0.invMassNSigma(Index), weight);
+            if (masterConfigurations.doMassSpectrumCheck) {
+              histos.fill(HIST("h") + HIST(V0names[Index]) + HIST("PtVsMass"), v0Data.pt(), getV0InvariantMass<Index>(v0Data));
+            }
             if (std::abs(v0Data.rapidity(Index)) < ySel) {
               histos.fill(HIST("h3d") + HIST(V0names[Index]) + HIST("SpectrumY"), v0Data.pt(), cent, v0.invMassNSigma(Index), weight);
             }