From e0e5202b913c23aa884114d2577903025b1a8bcf Mon Sep 17 00:00:00 2001
From: Fabio Catalano <fabio.catalano@cern.ch>
Date: Thu, 10 Oct 2024 16:10:12 +0200
Subject: [PATCH] Add trigger mass cut to Ds selector

---
 PWGHF/TableProducer/candidateSelectorDsToKKPi.cxx | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/PWGHF/TableProducer/candidateSelectorDsToKKPi.cxx b/PWGHF/TableProducer/candidateSelectorDsToKKPi.cxx
index 9720f270469..91df36d441c 100644
--- a/PWGHF/TableProducer/candidateSelectorDsToKKPi.cxx
+++ b/PWGHF/TableProducer/candidateSelectorDsToKKPi.cxx
@@ -25,6 +25,7 @@
 #include "PWGHF/Core/HfMlResponseDsToKKPi.h"
 #include "PWGHF/DataModel/CandidateReconstructionTables.h"
 #include "PWGHF/DataModel/CandidateSelectionTables.h"
+#include "PWGHF/Utils/utilsAnalysis.h"
 
 using namespace o2;
 using namespace o2::analysis;
@@ -73,6 +74,8 @@ struct HfCandidateSelectorDsToKKPi {
   Configurable<std::vector<std::string>> onnxFileNames{"onnxFileNames", std::vector<std::string>{"ModelHandler_onnx_DsToKKPi.onnx"}, "ONNX file names for each pT bin (if not from CCDB full path)"};
   Configurable<int64_t> timestampCCDB{"timestampCCDB", -1, "timestamp of the ONNX file for ML model used to query in CCDB"};
   Configurable<bool> loadModelsFromCCDB{"loadModelsFromCCDB", false, "Flag to enable or disable the loading of models from CCDB"};
+  // Mass cut for trigger analysis
+  Configurable<bool> useTriggerMassCut{"useTriggerMassCut", false, "Flag to enable parametrized pT differential mass cut for triggered data"};
 
   HfHelper hfHelper;
   o2::analysis::HfMlResponseDsToKKPi<float> hfMlResponse;
@@ -81,6 +84,7 @@ struct HfCandidateSelectorDsToKKPi {
   o2::ccdb::CcdbApi ccdbApi;
   TrackSelectorPi selectorPion;
   TrackSelectorKa selectorKaon;
+  HfTrigger3ProngCuts hfTriggerCuts;
 
   using TracksSel = soa::Join<aod::TracksWExtra, aod::TracksPidPi, aod::PidTpcTofFullPi, aod::TracksPidKa, aod::PidTpcTofFullKa>;
 
@@ -200,6 +204,9 @@ struct HfCandidateSelectorDsToKKPi {
     if (std::abs(hfHelper.invMassDsToKKPi(candidate) - o2::constants::physics::MassDS) > cuts->get(pTBin, "deltaM")) {
       return false;
     }
+    if (useTriggerMassCut && !isCandidateInMassRange(hfHelper.invMassDsToKKPi(candidate), o2::constants::physics::MassDS, candidate.pt(), hfTriggerCuts)) {
+      return false;
+    }
     if (hfHelper.deltaMassPhiDsToKKPi(candidate) > cuts->get(pTBin, "deltaM Phi")) {
       return false;
     }
@@ -229,6 +236,9 @@ struct HfCandidateSelectorDsToKKPi {
     if (std::abs(hfHelper.invMassDsToPiKK(candidate) - o2::constants::physics::MassDS) > cuts->get(pTBin, "deltaM")) {
       return false;
     }
+    if (useTriggerMassCut && !isCandidateInMassRange(hfHelper.invMassDsToPiKK(candidate), o2::constants::physics::MassDS, candidate.pt(), hfTriggerCuts)) {
+      return false;
+    }
     if (hfHelper.deltaMassPhiDsToPiKK(candidate) > cuts->get(pTBin, "deltaM Phi")) {
       return false;
     }