[PWGCF] a new table with ML score (#9850)

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

Author: jikim1290

PWGCF/DataModel/CorrelationsDerived.h

@@ -11,6 +11,8 @@
 #ifndef PWGCF_DATAMODEL_CORRELATIONSDERIVED_H_
 #define PWGCF_DATAMODEL_CORRELATIONSDERIVED_H_
 
+#include <vector>
+
 #include "Framework/ASoA.h"
 #include "Framework/AnalysisDataModel.h"
 #include "Common/DataModel/Centrality.h"
@@ -117,6 +119,20 @@ DECLARE_SOA_TABLE(CF2ProngTracks, "AOD", "CF2PRONGTRACK", //! Reduced track tabl
                   cf2prongtrack::CFTrackProng1Id,
                   cf2prongtrack::Pt, cf2prongtrack::Eta, cf2prongtrack::Phi, cf2prongtrack::InvMass, cf2prongtrack::Decay);
 using CF2ProngTrack = CF2ProngTracks::iterator;
+//------
+
+namespace cf2prongtrackml
+{
+DECLARE_SOA_COLUMN(MlProbD0, mlProbD0, std::vector<float>);       //!
+DECLARE_SOA_COLUMN(MlProbD0bar, mlProbD0bar, std::vector<float>); //!
+} // namespace cf2prongtrackml
+DECLARE_SOA_TABLE(CF2ProngTrackmls, "AOD", "CF2PRONGTRACKML", //! Reduced track table
+                  o2::soa::Index<>,
+                  cftrack::CFCollisionId,
+                  cf2prongtrackml::MlProbD0, cf2prongtrackml::MlProbD0bar);
+using CF2ProngTrackml = CF2ProngTrackmls::iterator;
+//------
+
 } // namespace o2::aod
 
 #endif // PWGCF_DATAMODEL_CORRELATIONSDERIVED_H_

PWGCF/TableProducer/filter2Prong.cxx

@@ -8,6 +8,8 @@
 // In applying this license CERN does not waive the privileges and immunities
 // granted to it by virtue of its status as an Intergovernmental Organization
 // or submit itself to any jurisdiction.
+#include <vector>
+
 #include "Framework/runDataProcessing.h"
 #include "Framework/AnalysisTask.h"
 #include "Framework/AnalysisDataModel.h"
@@ -35,8 +37,70 @@ struct Filter2Prong {
 
   HfHelper hfHelper;
   Produces<aod::CF2ProngTracks> output2ProngTracks;
+  Produces<aod::CF2ProngTrackmls> output2ProngTrackmls;
 
   using HFCandidates = soa::Join<aod::HfCand2Prong, aod::HfSelD0>;
+  using HFCandidatesML = soa::Join<aod::HfCand2Prong, aod::HfSelD0, aod::HfMlD0>;
+
+  void processDataML(aod::Collisions::iterator const&, aod::BCsWithTimestamps const&, aod::CFCollRefs const& cfcollisions, aod::CFTrackRefs const& cftracks, HFCandidatesML const& candidates)
+  {
+    if (cfcollisions.size() <= 0 || cftracks.size() <= 0)
+      return; // rejected collision
+    if (cfgVerbosity > 0 && candidates.size() > 0)
+      LOGF(info, "Candidates for collision: %lu, cfcollisions: %lu, CFTracks: %lu", candidates.size(), cfcollisions.size(), cftracks.size());
+    for (auto& c : candidates) {
+      int prongCFId[2] = {-1, -1};
+      for (auto& cftrack : cftracks) {
+        if (c.prong0Id() == cftrack.trackId()) {
+          prongCFId[0] = cftrack.globalIndex();
+          break;
+        }
+      }
+      for (auto& cftrack : cftracks) {
+        if (c.prong1Id() == cftrack.trackId()) {
+          prongCFId[1] = cftrack.globalIndex();
+          break;
+        }
+      }
+      // look-up the collision id
+      if ((c.hfflag() & (1 << aod::hf_cand_2prong::DecayType::D0ToPiK)) == 0)
+        continue;
+      if (cfgYMax >= 0.0f && std::abs(hfHelper.yD0(c)) > cfgYMax)
+        continue;
+
+      std::vector<float> mlvecd{};
+      std::vector<float> mlvecdbar{};
+
+      if (c.isSelD0() > 0) {
+        output2ProngTracks(cfcollisions.begin().globalIndex(),
+                           prongCFId[0], prongCFId[1], c.pt(), c.eta(), c.phi(), hfHelper.invMassD0ToPiK(c), aod::cf2prongtrack::D0ToPiK);
+        for (float val : c.mlProbD0()) {
+          mlvecd.push_back(val);
+        }
+        for (float val : c.mlProbD0bar()) {
+          mlvecdbar.push_back(val);
+        }
+        output2ProngTrackmls(cfcollisions.begin().globalIndex(), mlvecd, mlvecdbar);
+      }
+
+      mlvecd.clear();
+      mlvecdbar.clear();
+
+      if (c.isSelD0bar() > 0) {
+        output2ProngTracks(cfcollisions.begin().globalIndex(),
+                           prongCFId[0], prongCFId[1], c.pt(), c.eta(), c.phi(), hfHelper.invMassD0barToKPi(c), aod::cf2prongtrack::D0barToKPi);
+        for (float val : c.mlProbD0()) {
+          mlvecd.push_back(val);
+        }
+        for (float val : c.mlProbD0bar()) {
+          mlvecdbar.push_back(val);
+        }
+        output2ProngTrackmls(cfcollisions.begin().globalIndex(), mlvecd, mlvecdbar);
+      }
+    }
+  }
+  PROCESS_SWITCH(Filter2Prong, processDataML, "Process data D0 candidates with ML", false);
+
   void processData(aod::Collisions::iterator const&, aod::BCsWithTimestamps const&, aod::CFCollRefs const& cfcollisions, aod::CFTrackRefs const& cftracks, HFCandidates const& candidates)
   {
     if (cfcollisions.size() <= 0 || cftracks.size() <= 0)