[PWGJE] add enum D0McCategory to jetCorrelationD0

PWGJE/Tasks/jetCorrelationD0.cxx

@@ -13,6 +13,7 @@
 /// \brief Task for analysing D0 triggered jet events.
 /// \author Matthew Ockleton matthew.ockleton@cern.ch, University of Liverpool
 
+#include "PWGHF/Core/DecayChannels.h"
 #include "PWGJE/Core/JetDerivedDataUtilities.h"
 #include "PWGJE/Core/JetHFUtilities.h"
 #include "PWGJE/DataModel/Jet.h"
@@ -33,17 +34,17 @@
 
 #include <Rtypes.h>
 
+#include <cmath>
 #include <cstdint>
 #include <cstdlib>
 #include <string>
 #include <vector>
 
-#include <math.h>
-
 using namespace o2;
 using namespace o2::framework;
 using namespace o2::framework::expressions;
 
+namespace hf2Prong = o2::hf_decay::hf_cand_2prong;
 namespace o2::aod
 {
 
@@ -86,8 +87,7 @@ DECLARE_SOA_COLUMN(D0MD, d0MD, float);
 DECLARE_SOA_COLUMN(D0PtD, d0PtD, float);
 DECLARE_SOA_COLUMN(D0EtaD, d0EtaD, float);
 DECLARE_SOA_COLUMN(D0PhiD, d0PhiD, float);
-DECLARE_SOA_COLUMN(D0MatchedFrom, d0MatchedFrom, int);
-DECLARE_SOA_COLUMN(D0SelectedAs, d0SelectedAs, int);
+DECLARE_SOA_COLUMN(D0Category, d0Category, int);
 DECLARE_SOA_COLUMN(D0DecayChannel, d0DecayChannel, int8_t);
 } // namespace d0Info
 
@@ -114,9 +114,7 @@ DECLARE_SOA_TABLE(D0McDTables, "AOD", "D0MCDTABLE",
                   d0Info::D0Eta,
                   d0Info::D0Phi,
                   d0Info::D0Y,
-                  d0Info::D0MatchedFrom,
-                  d0Info::D0SelectedAs,
-                  d0Info::D0DecayChannel);
+                  d0Info::D0Category);
 
 DECLARE_SOA_TABLE(D0McPTables, "AOD", "D0MCPTABLE",
                   o2::soa::Index<>,
@@ -293,6 +291,15 @@ struct JetCorrelationD0 {
     registry.add("hPhiResolution", "#phi resolution;#p_{T,part};Resolution", {HistType::kTH2F, {{400, 0, 400}, {1000, -7.0, 7.0}}});
     registry.add("hEtaResolution", "#eta resolution;#p_{T,part};Resolution", {HistType::kTH2F, {{400, 0, 400}, {1000, -1.0, 1.0}}});
   }
+  enum D0McCategory : int {
+    Undefined = -1,     // no truth match / unclassified
+    Signal = 0,         // correctly identified D0(bar), π+ K−
+    Reflection = 1,     // true D0(bar) reconstructed with swapped mass hypothesis
+    CorrBkgPiKPi0 = 2,  // correlated background: π+ K− π0
+    CorrBkgPiPi = 3,    // correlated background: π+ π−
+    CorrBkgPiPiPi0 = 4, // correlated background: π+ π− π0
+    CorrBkgKK = 5       // correlated background: K+ K−
+  };
   void processData(soa::Filtered<aod::JetCollisions>::iterator const& collision,
                    aod::CandidatesD0Data const& d0Candidates,
                    soa::Join<aod::ChargedJets, aod::ChargedJetConstituents> const& jets)
@@ -357,17 +364,31 @@ struct JetCorrelationD0 {
 
       int matchedFrom = 0;
       int selectedAs = 0;
+      int category = D0McCategory::Undefined;
 
       if (d0DecayChannel > 0) { // matched to a D0 on truth level (any channel)
         matchedFrom = 1;
       } else if (d0DecayChannel < 0) { // matched to a D0bar on truth level (any channel)
         matchedFrom = -1;
       }
-      if (d0Candidate.candidateSelFlag() & BIT(0)) { // CandidateSelFlag == BIT(0) -> selected as D0
+      if ((d0Candidate.candidateSelFlag() & BIT(0)) != 0) { // CandidateSelFlag == BIT(0) -> selected as D0
         selectedAs = 1;
-      } else if (d0Candidate.candidateSelFlag() & BIT(1)) { // CandidateSelFlag == BIT(1) -> selected as D0bar
+      } else if ((d0Candidate.candidateSelFlag() & BIT(1)) != 0) { // CandidateSelFlag == BIT(1) -> selected as D0bar
         selectedAs = -1;
       }
+      if ((std::abs(d0DecayChannel) == hf2Prong::DecayChannelMain::D0ToPiK) && (matchedFrom != 0) && (selectedAs == matchedFrom)) {
+        category = D0McCategory::Signal; // D0 or D0bar, π+ K−
+      } else if ((d0DecayChannel == hf2Prong::DecayChannelMain::D0ToPiK) && (selectedAs == -1 * matchedFrom)) {
+        category = D0McCategory::Reflection;
+      } else if (d0DecayChannel == hf2Prong::DecayChannelMain::D0ToPiKPi0) {
+        category = D0McCategory::CorrBkgPiKPi0;
+      } else if (d0DecayChannel == hf2Prong::DecayChannelMain::D0ToPiPi) {
+        category = D0McCategory::CorrBkgPiPi;
+      } else if (d0DecayChannel == hf2Prong::DecayChannelMain::D0ToPiPiPi0) {
+        category = D0McCategory::CorrBkgPiPiPi0;
+      } else if (d0DecayChannel == hf2Prong::DecayChannelMain::D0ToKK) {
+        category = D0McCategory::CorrBkgKK;
+      }
 
       tableD0McDetector(tableCollision.lastIndex(), // might want to add some more detector level D0 quantities like prompt or non prompt info
                         scores[2],
@@ -378,9 +399,7 @@ struct JetCorrelationD0 {
                         d0Candidate.eta(),
                         d0Candidate.phi(),
                         d0Candidate.y(),
-                        matchedFrom,
-                        selectedAs,
-                        d0DecayChannel);
+                        category);
       for (const auto& jet : jets) {
         if (jet.pt() < jetPtCutMin) {
           continue;