[PWGHF] Add table with ML outputs for reco MC (#10063)

Co-authored-by: ALICE Action Bot <alibuild@cern.ch>

Author: Marcellocosti

PWGHF/DataModel/CandidateReconstructionTables.h

@@ -949,6 +949,8 @@ enum DecayType { DplusToPiKPi = 0,
                  XicToPKPi,
                  N3ProngDecays }; // always keep N3ProngDecays at the end
 
+static const int DstarToPiKPiBkg = DecayType::N3ProngDecays;
+
 // Ds± → K± K∓ π± or D± → K± K∓ π±
 
 enum DecayChannelDToKKPi {

PWGHF/TableProducer/candidateCreator3Prong.cxx

@@ -787,11 +787,6 @@ struct HfCandidateCreator3ProngExpressions {
   o2::framework::Configurable<bool> matchKinkedDecayTopology{"matchKinkedDecayTopology", false, "Match also candidates with tracks that decay with kinked topology"};
   o2::framework::Configurable<bool> matchInteractionsWithMaterial{"matchInteractionsWithMaterial", false, "Match also candidates with tracks that interact with material"};
 
-  bool createDplus{false};
-  bool createDs{false};
-  bool createLc{false};
-  bool createXic{false};
-
   HfEventSelectionMc hfEvSelMc; // mc event selection and monitoring
   using BCsInfo = soa::Join<aod::BCs, aod::Timestamps, aod::BcSels>;
   HistogramRegistry registry{"registry"};
@@ -817,28 +812,11 @@ struct HfCandidateCreator3ProngExpressions {
     for (const DeviceSpec& device : workflows.devices) {
       if (device.name.compare("hf-candidate-creator-3prong") == 0) {
         hfEvSelMc.configureFromDevice(device);
-        for (const auto& option : device.options) {
-          if (option.name.compare("createDplus") == 0) {
-            createDplus = option.defaultValue.get<bool>();
-          } else if (option.name.compare("createDs") == 0) {
-            createDs = option.defaultValue.get<bool>();
-          } else if (option.name.compare("createLc") == 0) {
-            createLc = option.defaultValue.get<bool>();
-          } else if (option.name.compare("createXic") == 0) {
-            createXic = option.defaultValue.get<bool>();
-          }
-        }
         break;
       }
     }
 
     hfEvSelMc.addHistograms(registry); // particles monitoring
-
-    LOGP(info, "Flags for candidate creation from the reco workflow:");
-    LOGP(info, "    --> createDplus = {}", createDplus);
-    LOGP(info, "    --> createDs = {}", createDs);
-    LOGP(info, "    --> createLc = {}", createLc);
-    LOGP(info, "    --> createXic = {}", createXic);
   }
 
   /// Performs MC matching.
@@ -897,7 +875,7 @@ struct HfCandidateCreator3ProngExpressions {
       }
 
       // D± → π± K∓ π±
-      if (createDplus) {
+      if (flag == 0) {
         if (matchKinkedDecayTopology && matchInteractionsWithMaterial) {
           indexRec = RecoDecay::getMatchedMCRec<false, false, false, true, true>(mcParticles, arrayDaughters, Pdg::kDPlus, std::array{+kPiPlus, -kKPlus, +kPiPlus}, true, &sign, 2, &nKinkedTracks, &nInteractionsWithMaterial);
         } else if (matchKinkedDecayTopology && !matchInteractionsWithMaterial) {
@@ -913,7 +891,7 @@ struct HfCandidateCreator3ProngExpressions {
       }
 
       // Ds± → K± K∓ π± and D± → K± K∓ π±
-      if (flag == 0 && createDs) {
+      if (flag == 0) {
         bool isDplus = false;
         if (matchKinkedDecayTopology && matchInteractionsWithMaterial) {
           indexRec = RecoDecay::getMatchedMCRec<false, false, false, true, true>(mcParticles, arrayDaughters, Pdg::kDS, std::array{+kKPlus, -kKPlus, +kPiPlus}, true, &sign, 2, &nKinkedTracks, &nInteractionsWithMaterial);
@@ -958,8 +936,21 @@ struct HfCandidateCreator3ProngExpressions {
         }
       }
 
+      // D* → D0π → Kππ
+      if (flag == 0) {
+        if (matchKinkedDecayTopology) {
+          indexRec = RecoDecay::getMatchedMCRec<false, false, false, true>(mcParticles, arrayDaughters, Pdg::kDStar, std::array{+kPiPlus, +kPiPlus, -kKPlus}, true, &sign, 2, &nKinkedTracks);
+        } else {
+          indexRec = RecoDecay::getMatchedMCRec(mcParticles, arrayDaughters, Pdg::kDStar, std::array{+kPiPlus, +kPiPlus, -kKPlus}, true, &sign, 2);
+        }
+        if (indexRec > -1) {
+          flag = sign * (1 << DstarToPiKPiBkg);
+          channel = 1;
+        }
+      }
+
       // Λc± → p± K∓ π±
-      if (flag == 0 && createLc) {
+      if (flag == 0) {
         if (matchKinkedDecayTopology && matchInteractionsWithMaterial) {
           indexRec = RecoDecay::getMatchedMCRec<false, false, false, true, true>(mcParticles, arrayDaughters, Pdg::kLambdaCPlus, std::array{+kProton, -kKPlus, +kPiPlus}, true, &sign, 2, &nKinkedTracks, &nInteractionsWithMaterial);
         } else if (matchKinkedDecayTopology && !matchInteractionsWithMaterial) {
@@ -994,7 +985,7 @@ struct HfCandidateCreator3ProngExpressions {
       }
 
       // Ξc± → p± K∓ π±
-      if (flag == 0 && createXic) {
+      if (flag == 0) {
         if (matchKinkedDecayTopology && matchInteractionsWithMaterial) {
           indexRec = RecoDecay::getMatchedMCRec<false, false, false, true, true>(mcParticles, arrayDaughters, Pdg::kXiCPlus, std::array{+kProton, -kKPlus, +kPiPlus}, true, &sign, 2, &nKinkedTracks, &nInteractionsWithMaterial);
         } else if (matchKinkedDecayTopology && !matchInteractionsWithMaterial) {
@@ -1049,7 +1040,7 @@ struct HfCandidateCreator3ProngExpressions {
         }
         continue;
       }
-      hf_mc_gen::fillMcMatchGen3Prong(mcParticles, mcParticlesPerMcColl, rowMcMatchGen, rejectBackground, createDplus, createDs, createLc, createXic);
+      hf_mc_gen::fillMcMatchGen3Prong(mcParticles, mcParticlesPerMcColl, rowMcMatchGen, rejectBackground);
     }
   }
 

PWGHF/TableProducer/candidateCreatorMcGen.cxx

@@ -46,10 +46,6 @@ struct HfCandidateCreatorMcGen {
   Configurable<bool> fillB0{"fillB0", false, "fill table for B0 candidates"};
   Configurable<bool> rejectBackground2Prong{"rejectBackground2Prong", false, "Reject particles from PbPb background for 2 prong candidates"};
   Configurable<bool> rejectBackground3Prong{"rejectBackground3Prong", false, "Reject particles from PbPb background for 3 prong candidates"};
-  Configurable<bool> createDplus{"createDplus", false, "Create D+ in 3 prong"};
-  Configurable<bool> createDs{"createDs", false, "Create Ds in 3 prong"};
-  Configurable<bool> createLc{"createLc", false, "Create Lc in 3 prong"};
-  Configurable<bool> createXic{"createXic", false, "Create Xic in 3 prong"};
 
   Preslice<aod::McParticles> mcParticlesPerMcCollision = aod::mcparticle::mcCollisionId;
 
@@ -63,7 +59,7 @@ struct HfCandidateCreatorMcGen {
         hf_mc_gen::fillMcMatchGen2Prong(mcParticles, mcParticlesPerMcColl, rowMcMatchGen2Prong, rejectBackground2Prong);
       }
       if (fill3Prong) {
-        hf_mc_gen::fillMcMatchGen3Prong(mcParticles, mcParticlesPerMcColl, rowMcMatchGen3Prong, rejectBackground3Prong, createDplus, createDs, createLc, createXic);
+        hf_mc_gen::fillMcMatchGen3Prong(mcParticles, mcParticlesPerMcColl, rowMcMatchGen3Prong, rejectBackground3Prong);
       }
     }
     if (fillBplus) {

PWGHF/TableProducer/treeCreatorDplusToPiKPi.cxx

@@ -79,7 +79,13 @@ DECLARE_SOA_COLUMN(Ct, ct, float);
 // Events
 DECLARE_SOA_COLUMN(IsEventReject, isEventReject, int); //! Event rejection flag
 DECLARE_SOA_COLUMN(RunNumber, runNumber, int);         //! Run number
+// ML scores
+DECLARE_SOA_COLUMN(MlScore0, mlScore0, float); //! ML score of the first configured index
+DECLARE_SOA_COLUMN(MlScore1, mlScore1, float); //! ML score of the second configured index
 } // namespace full
+DECLARE_SOA_TABLE(HfCandDpMls, "AOD", "HFCANDDPML",
+                  full::MlScore0,
+                  full::MlScore1)
 
 DECLARE_SOA_TABLE(HfCandDpLites, "AOD", "HFCANDDPLITE",
                   hf_cand::Chi2PCA,
@@ -233,26 +239,31 @@ struct HfTreeCreatorDplusToPiKPi {
   Produces<o2::aod::HfCandDpFullEvs> rowCandidateFullEvents;
   Produces<o2::aod::HfCandDpFullPs> rowCandidateFullParticles;
   Produces<o2::aod::HfCandDpLites> rowCandidateLite;
+  Produces<o2::aod::HfCandDpMls> rowCandidateMl;
 
   Configurable<int> selectionFlagDplus{"selectionFlagDplus", 1, "Selection Flag for Dplus"};
   Configurable<bool> fillCandidateLiteTable{"fillCandidateLiteTable", false, "Switch to fill lite table with candidate properties"};
   // parameters for production of training samples
   Configurable<bool> fillOnlySignal{"fillOnlySignal", false, "Flag to fill derived tables with signal for ML trainings"};
+  Configurable<bool> fillOnlySignalMl{"fillOnlySignalMl", false, "Flag to fill derived tables with MC and ML info"};
   Configurable<bool> fillOnlyBackground{"fillOnlyBackground", false, "Flag to fill derived tables with background for ML trainings"};
   Configurable<float> downSampleBkgFactor{"downSampleBkgFactor", 1., "Fraction of background candidates to keep for ML trainings"};
   Configurable<float> ptMaxForDownSample{"ptMaxForDownSample", 10., "Maximum pt for the application of the downsampling factor"};
+  Configurable<std::vector<int>> classMl{"classMlindexes", {0, 2}, "Indexes of ML bkg and non-prompt scores."};
 
   HfHelper hfHelper;
 
   using SelectedCandidatesMc = soa::Filtered<soa::Join<aod::HfCand3Prong, aod::HfCand3ProngMcRec, aod::HfSelDplusToPiKPi>>;
   using MatchedGenCandidatesMc = soa::Filtered<soa::Join<aod::McParticles, aod::HfCand3ProngMcGen>>;
+  using SelectedCandidatesMcWithMl = soa::Filtered<soa::Join<aod::HfCand3Prong, aod::HfCand3ProngMcRec, aod::HfSelDplusToPiKPi, aod::HfMlDplusToPiKPi>>;
   using TracksWPid = soa::Join<aod::Tracks, aod::TracksPidPi, aod::PidTpcTofFullPi, aod::TracksPidKa, aod::PidTpcTofFullKa>;
 
   Filter filterSelectCandidates = aod::hf_sel_candidate_dplus::isSelDplusToPiKPi >= selectionFlagDplus;
   Filter filterMcGenMatching = nabs(o2::aod::hf_cand_3prong::flagMcMatchGen) == static_cast<int8_t>(BIT(aod::hf_cand_3prong::DecayType::DplusToPiKPi));
 
   Partition<SelectedCandidatesMc> reconstructedCandSig = nabs(aod::hf_cand_3prong::flagMcMatchRec) == static_cast<int8_t>(BIT(aod::hf_cand_3prong::DecayType::DplusToPiKPi)) || nabs(aod::hf_cand_3prong::flagMcMatchRec) == static_cast<int8_t>(BIT(aod::hf_cand_3prong::DecayType::DsToKKPi)); // DecayType::DsToKKPi is used to flag both Ds± → K± K∓ π± and D± → K± K∓ π±
   Partition<SelectedCandidatesMc> reconstructedCandBkg = nabs(aod::hf_cand_3prong::flagMcMatchRec) != static_cast<int8_t>(BIT(aod::hf_cand_3prong::DecayType::DplusToPiKPi));
+  Partition<SelectedCandidatesMcWithMl> reconstructedCandSigMl = nabs(aod::hf_cand_3prong::flagMcMatchRec) == static_cast<int8_t>(BIT(aod::hf_cand_3prong::DecayType::DplusToPiKPi)) || nabs(aod::hf_cand_3prong::flagMcMatchRec) == static_cast<int8_t>(BIT(aod::hf_cand_3prong::DecayType::DsToKKPi)) || nabs(aod::hf_cand_3prong::flagMcMatchRec) == static_cast<int8_t>(BIT(aod::hf_cand_3prong::DstarToPiKPiBkg)); // DecayType::DsToKKPi is used to flag both Ds± → K± K∓ π± and D± → K± K∓ π±
 
   void init(InitContext const&)
   {
@@ -271,7 +282,7 @@ struct HfTreeCreatorDplusToPiKPi {
       runNumber);
   }
 
-  template <bool doMc = false, typename T>
+  template <bool doMc = false, bool doMl = false, typename T>
   void fillCandidateTable(const T& candidate)
   {
     int8_t flagMc = 0;
@@ -283,6 +294,16 @@ struct HfTreeCreatorDplusToPiKPi {
       channelMc = candidate.flagMcDecayChanRec();
     }
 
+    std::vector<float> outputMl = {-999., -999.};
+    if constexpr (doMl) {
+      for (unsigned int iclass = 0; iclass < classMl->size(); iclass++) {
+        outputMl[iclass] = candidate.mlProbDplusToPiKPi()[classMl->at(iclass)];
+      }
+      rowCandidateMl(
+        outputMl[0],
+        outputMl[1]);
+    }
+
     auto prong0 = candidate.template prong0_as<TracksWPid>();
     auto prong1 = candidate.template prong1_as<TracksWPid>();
     auto prong2 = candidate.template prong2_as<TracksWPid>();
@@ -450,6 +471,7 @@ struct HfTreeCreatorDplusToPiKPi {
                  aod::McCollisions const&,
                  SelectedCandidatesMc const& candidates,
                  MatchedGenCandidatesMc const& particles,
+                 SelectedCandidatesMcWithMl const&,
                  TracksWPid const&)
   {
     // Filling event properties
@@ -468,6 +490,22 @@ struct HfTreeCreatorDplusToPiKPi {
       for (const auto& candidate : reconstructedCandSig) {
         fillCandidateTable<true>(candidate);
       }
+    } else if (fillOnlySignalMl) {
+      rowCandidateMl.reserve(reconstructedCandSigMl.size());
+      if (fillCandidateLiteTable) {
+        rowCandidateLite.reserve(reconstructedCandSigMl.size());
+      } else {
+        rowCandidateFull.reserve(reconstructedCandSigMl.size());
+      }
+      for (const auto& candidate : reconstructedCandSigMl) {
+        if (downSampleBkgFactor < 1.) {
+          float pseudoRndm = candidate.ptProng0() * 1000. - (int64_t)(candidate.ptProng0() * 1000);
+          if (candidate.pt() < ptMaxForDownSample && pseudoRndm >= downSampleBkgFactor) {
+            continue;
+          }
+        }
+        fillCandidateTable<true, true>(candidate);
+      }
     } else if (fillOnlyBackground) {
       if (fillCandidateLiteTable) {
         rowCandidateLite.reserve(reconstructedCandBkg.size());

PWGHF/Utils/utilsMcGen.h

@@ -79,7 +79,7 @@ void fillMcMatchGen2Prong(T const& mcParticles, U const& mcParticlesPerMcColl, V
 }
 
 template <typename T, typename U, typename V>
-void fillMcMatchGen3Prong(T const& mcParticles, U const& mcParticlesPerMcColl, V& rowMcMatchGen, bool rejectBackground, bool createDplus, bool createDs, bool createLc, bool createXic)
+void fillMcMatchGen3Prong(T const& mcParticles, U const& mcParticlesPerMcColl, V& rowMcMatchGen, bool rejectBackground)
 {
   using namespace o2::constants::physics;
 
@@ -104,14 +104,14 @@ void fillMcMatchGen3Prong(T const& mcParticles, U const& mcParticlesPerMcColl, V
     }
 
     // D± → π± K∓ π±
-    if (createDplus) {
+    if (flag == 0) {
       if (RecoDecay::isMatchedMCGen(mcParticles, particle, Pdg::kDPlus, std::array{+kPiPlus, -kKPlus, +kPiPlus}, true, &sign, 2)) {
         flag = sign * (1 << o2::aod::hf_cand_3prong::DecayType::DplusToPiKPi);
       }
     }
 
     // Ds± → K± K∓ π± and D± → K± K∓ π±
-    if (flag == 0 && createDs) {
+    if (flag == 0) {
       bool isDplus = false;
       if (RecoDecay::isMatchedMCGen(mcParticles, particle, Pdg::kDS, std::array{+kKPlus, -kKPlus, +kPiPlus}, true, &sign, 2)) {
         // DecayType::DsToKKPi is used to flag both Ds± → K± K∓ π± and D± → K± K∓ π±
@@ -139,8 +139,15 @@ void fillMcMatchGen3Prong(T const& mcParticles, U const& mcParticlesPerMcColl, V
       }
     }
 
+    // D*± → D0(bar) π±
+    if (flag == 0) {
+      if (RecoDecay::isMatchedMCGen(mcParticles, particle, Pdg::kDStar, std::array{+kPiPlus, +kPiPlus, -kKPlus}, true, &sign, 2)) {
+        flag = sign * (1 << o2::aod::hf_cand_3prong::DstarToPiKPiBkg);
+      }
+    }
+
     // Λc± → p± K∓ π±
-    if (flag == 0 && createLc) {
+    if (flag == 0) {
       if (RecoDecay::isMatchedMCGen(mcParticles, particle, Pdg::kLambdaCPlus, std::array{+kProton, -kKPlus, +kPiPlus}, true, &sign, 2)) {
         flag = sign * (1 << o2::aod::hf_cand_3prong::DecayType::LcToPKPi);
 
@@ -163,7 +170,7 @@ void fillMcMatchGen3Prong(T const& mcParticles, U const& mcParticlesPerMcColl, V
     }
 
     // Ξc± → p± K∓ π±
-    if (flag == 0 && createXic) {
+    if (flag == 0) {
       if (RecoDecay::isMatchedMCGen(mcParticles, particle, Pdg::kXiCPlus, std::array{+kProton, -kKPlus, +kPiPlus}, true, &sign, 2)) {
         flag = sign * (1 << o2::aod::hf_cand_3prong::DecayType::XicToPKPi);
       }