[PWGHF] Fix ML variable addition to tree creator (AliceO2Group#9885)

Author: atavirag

PWGHF/TableProducer/treeCreatorD0ToKPi.cxx

@@ -17,6 +17,8 @@
 /// \author Nicolo' Jacazio <[email protected]>, CERN
 /// \author Andrea Tavira García <[email protected]>, IJCLab
 
+#include <vector>
+
 #include "CommonConstants/PhysicsConstants.h"
 #include "Framework/AnalysisTask.h"
 #include "Framework/runDataProcessing.h"
@@ -275,7 +277,7 @@ struct HfTreeCreatorD0ToKPi {
 
   template <bool applyMl, typename T>
   auto fillTable(const T& candidate, int candFlag, double invMass, double cosThetaStar, double topoChi2,
-                 double ct, double y, double e, int8_t flagMc, int8_t origin)
+                 double ct, double y, double e, int8_t flagMc, int8_t origin, std::vector<float> const& mlProbD0)
   {
     if (fillCandidateLiteTable) {
       rowCandidateLite(
@@ -379,18 +381,10 @@ struct HfTreeCreatorD0ToKPi {
         candidate.globalIndex());
     }
     if constexpr (applyMl) {
-      if (candidate.isSelD0()) {
-        rowCandidateMl(
-          candidate.mlProbD0()[0],
-          candidate.mlProbD0()[1],
-          candidate.mlProbD0()[2]);
-      }
-      if (candidate.isSelD0bar()) {
-        rowCandidateMl(
-          candidate.mlProbD0bar()[0],
-          candidate.mlProbD0bar()[1],
-          candidate.mlProbD0bar()[2]);
-      }
+      rowCandidateMl(
+        mlProbD0[0],
+        mlProbD0[1],
+        mlProbD0[2]);
     }
   }
 
@@ -425,6 +419,7 @@ struct HfTreeCreatorD0ToKPi {
       double eD = hfHelper.eD0(candidate);
       double ctD = hfHelper.ctD0(candidate);
       float massD0, massD0bar;
+      std::vector<float> mlProbD0, mlProbD0bar;
       float topolChi2PerNdf = -999.;
       if constexpr (reconstructionType == aod::hf_cand::VertexerType::KfParticle) {
         massD0 = candidate.kfGeoMassD0();
@@ -435,10 +430,16 @@ struct HfTreeCreatorD0ToKPi {
         massD0bar = hfHelper.invMassD0barToKPi(candidate);
       }
       if (candidate.isSelD0()) {
-        fillTable<applyMl>(candidate, 0, massD0, hfHelper.cosThetaStarD0(candidate), topolChi2PerNdf, ctD, yD, eD, 0, 0);
+        if constexpr (applyMl) {
+          mlProbD0 = std::vector<float>(candidate.mlProbD0().begin(), candidate.mlProbD0().end());
+        }
+        fillTable<applyMl>(candidate, 0, massD0, hfHelper.cosThetaStarD0(candidate), topolChi2PerNdf, ctD, yD, eD, 0, 0, mlProbD0);
       }
       if (candidate.isSelD0bar()) {
-        fillTable<applyMl>(candidate, 1, massD0bar, hfHelper.cosThetaStarD0bar(candidate), topolChi2PerNdf, ctD, yD, eD, 0, 0);
+        if constexpr (applyMl) {
+          mlProbD0bar = std::vector<float>(candidate.mlProbD0bar().begin(), candidate.mlProbD0bar().end());
+        }
+        fillTable<applyMl>(candidate, 1, massD0bar, hfHelper.cosThetaStarD0bar(candidate), topolChi2PerNdf, ctD, yD, eD, 0, 0, mlProbD0bar);
       }
     }
   }
@@ -524,6 +525,7 @@ struct HfTreeCreatorD0ToKPi {
       double ctD = hfHelper.ctD0(candidate);
       float massD0, massD0bar;
       float topolChi2PerNdf = -999.;
+      std::vector<float> mlProbD0, mlProbD0bar;
       if constexpr (reconstructionType == aod::hf_cand::VertexerType::KfParticle) {
         massD0 = candidate.kfGeoMassD0();
         massD0bar = candidate.kfGeoMassD0bar();
@@ -533,10 +535,16 @@ struct HfTreeCreatorD0ToKPi {
         massD0bar = hfHelper.invMassD0barToKPi(candidate);
       }
       if (candidate.isSelD0()) {
-        fillTable<applyMl>(candidate, 0, massD0, hfHelper.cosThetaStarD0(candidate), topolChi2PerNdf, ctD, yD, eD, candidate.flagMcMatchRec(), candidate.originMcRec());
+        if constexpr (applyMl) {
+          mlProbD0 = std::vector<float>(candidate.mlProbD0().begin(), candidate.mlProbD0().end());
+        }
+        fillTable<applyMl>(candidate, 0, massD0, hfHelper.cosThetaStarD0(candidate), topolChi2PerNdf, ctD, yD, eD, candidate.flagMcMatchRec(), candidate.originMcRec(), mlProbD0);
       }
       if (candidate.isSelD0bar()) {
-        fillTable<applyMl>(candidate, 1, massD0bar, hfHelper.cosThetaStarD0bar(candidate), topolChi2PerNdf, ctD, yD, eD, candidate.flagMcMatchRec(), candidate.originMcRec());
+        if constexpr (applyMl) {
+          mlProbD0bar = std::vector<float>(candidate.mlProbD0bar().begin(), candidate.mlProbD0bar().end());
+        }
+        fillTable<applyMl>(candidate, 1, massD0bar, hfHelper.cosThetaStarD0bar(candidate), topolChi2PerNdf, ctD, yD, eD, candidate.flagMcMatchRec(), candidate.originMcRec(), mlProbD0bar);
       }
     }