[PWGHF] Cleanup table filling in D0 tree creator

Author: fcatalan92

PWGHF/TableProducer/treeCreatorD0ToKPi.cxx

@@ -276,8 +276,8 @@ 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, std::vector<float> const& mlProbD0)
+  auto fillTable(const T& candidate, int candFlag, double invMass, double topoChi2,
+                 double ct, double y, double e, int8_t flagMc, int8_t origin)
   {
     if (fillCandidateLiteTable) {
       rowCandidateLite(
@@ -317,6 +317,7 @@ struct HfTreeCreatorD0ToKPi {
         flagMc,
         origin);
     } else {
+      double cosThetaStar = candFlag == 0 ? hfHelper.cosThetaStarD0(candidate) : hfHelper.cosThetaStarD0bar(candidate);
       rowCandidateFull(
         candidate.collisionId(),
         candidate.posX(),
@@ -381,10 +382,17 @@ struct HfTreeCreatorD0ToKPi {
         candidate.globalIndex());
     }
     if constexpr (applyMl) {
-      rowCandidateMl(
-        mlProbD0[0],
-        mlProbD0[1],
-        mlProbD0[2]);
+      if (candFlag == 0) {
+        rowCandidateMl(
+          candidate.mlProbD0()[0],
+          candidate.mlProbD0()[1],
+          candidate.mlProbD0()[2]);
+      } else if (candFlag == 1) {
+        rowCandidateMl(
+          candidate.mlProbD0bar()[0],
+          candidate.mlProbD0bar()[1],
+          candidate.mlProbD0bar()[2]);
+      }
     }
   }
 
@@ -419,7 +427,6 @@ 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();
@@ -430,16 +437,10 @@ struct HfTreeCreatorD0ToKPi {
         massD0bar = hfHelper.invMassD0barToKPi(candidate);
       }
       if (candidate.isSelD0()) {
-        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);
+        fillTable<applyMl>(candidate, 0, massD0, topolChi2PerNdf, ctD, yD, eD, 0, 0);
       }
       if (candidate.isSelD0bar()) {
-        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);
+        fillTable<applyMl>(candidate, 1, massD0bar, topolChi2PerNdf, ctD, yD, eD, 0, 0);
       }
     }
   }
@@ -525,7 +526,6 @@ 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();
@@ -535,16 +535,10 @@ struct HfTreeCreatorD0ToKPi {
         massD0bar = hfHelper.invMassD0barToKPi(candidate);
       }
       if (candidate.isSelD0()) {
-        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);
+        fillTable<applyMl>(candidate, 0, massD0, topolChi2PerNdf, ctD, yD, eD, candidate.flagMcMatchRec(), candidate.originMcRec());
       }
       if (candidate.isSelD0bar()) {
-        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);
+        fillTable<applyMl>(candidate, 1, massD0bar, topolChi2PerNdf, ctD, yD, eD, candidate.flagMcMatchRec(), candidate.originMcRec());
       }
     }