@@ -336,20 +336,29 @@ struct HfTaskD0 {
336336 axes.push_back (thnAxisMinItsNCls);
337337 axes.push_back (thnAxisMinTpcNCrossedRows);
338338 }
339- if (doprocessDataWithDCAFitterNMlWithUpc) {
340- axes.push_back (thnAxisGapType);
341- axes.push_back (thnAxisFT0A);
342- axes.push_back (thnAxisFT0C);
343- }
344339 if (applyMl) {
345340 const AxisSpec thnAxisBkgScore{thnConfigAxisBkgScore, " BDT score bkg." };
346341 const AxisSpec thnAxisNonPromptScore{thnConfigAxisNonPromptScore, " BDT score non-prompt." };
347342 const AxisSpec thnAxisPromptScore{thnConfigAxisPromptScore, " BDT score prompt." };
348343
349- axes.insert (axes.begin (), thnAxisPromptScore);
350- axes.insert (axes.begin (), thnAxisNonPromptScore);
351- axes.insert (axes.begin (), thnAxisBkgScore);
344+ // Insert ML scores after pt (position 2) to match taskDplus structure: [mass, pt, mlScores, ...]
345+ if (doprocessDataWithDCAFitterNMlWithUpc) {
346+ axes.insert (axes.begin () + 2 , thnAxisPromptScore);
347+ axes.insert (axes.begin () + 2 , thnAxisNonPromptScore);
348+ axes.insert (axes.begin () + 2 , thnAxisBkgScore);
349+ } else {
350+ axes.insert (axes.begin (), thnAxisPromptScore);
351+ axes.insert (axes.begin (), thnAxisNonPromptScore);
352+ axes.insert (axes.begin (), thnAxisBkgScore);
353+ }
354+ }
355+ if (doprocessDataWithDCAFitterNMlWithUpc) {
356+ axes.push_back (thnAxisGapType);
357+ axes.push_back (thnAxisFT0A);
358+ axes.push_back (thnAxisFT0C);
359+ }
352360
361+ if (applyMl) {
353362 registry.add (" hBdtScoreVsMassVsPtVsPtBVsYVsOriginVsD0Type" , " Thn for D0 candidates" , HistType::kTHnSparseD , axes);
354363 registry.get <THnSparse>(HIST (" hBdtScoreVsMassVsPtVsPtBVsYVsOriginVsD0Type" ))->Sumw2 ();
355364 } else {
@@ -582,6 +591,8 @@ struct HfTaskD0 {
582591 int const gapTypeInt = hf_upc::gapTypeToInt (gap);
583592 const auto thisCollId = collision.globalIndex ();
584593 const auto & groupedD0Candidates = candidates.sliceBy (candD0PerCollision, thisCollId);
594+ float cent{centrality};
595+ float occ{-1 .f };
585596
586597 for (const auto & candidate : groupedD0Candidates) {
587598 if (!(candidate.hfflag () & 1 << aod::hf_cand_2prong::DecayType::D0ToPiK)) {
@@ -606,12 +617,18 @@ struct HfTaskD0 {
606617 registry.fill (HIST (" hMassVsPhi" ), massD0bar, ptCandidate, candidate.phi ());
607618 }
608619
609- // Fill THnSparse with gap type and FIT signals using vectorized approach
620+ // Fill THnSparse with structure matching taskDplus: [mass, pt, mlScores, cent, occ, gapType, FT0A, FT0C]
610621 if constexpr (fillMl) {
611622 auto fillTHnData = [&](float mass, int d0Type) {
612- std::vector<double > valuesToFill{candidate. mlProbD0 ()[ 0 ], candidate. mlProbD0 ()[ 1 ], candidate. mlProbD0 ()[ 2 ] ,
613- static_cast < double >(mass), static_cast < double >(ptCandidate) ,
623+ std::vector<double > valuesToFill{static_cast < double >(mass), static_cast < double >(ptCandidate) ,
624+ candidate. mlProbD0 ()[ 0 ], candidate. mlProbD0 ()[ 1 ], candidate. mlProbD0 ()[ 2 ] ,
614625 static_cast <double >(HfHelper::yD0 (candidate)), static_cast <double >(d0Type)};
626+ if (storeCentrality) {
627+ valuesToFill.push_back (cent);
628+ }
629+ if (storeOccupancyAndIR) {
630+ valuesToFill.push_back (occ);
631+ }
615632 valuesToFill.push_back (static_cast <double >(gapTypeInt));
616633 valuesToFill.push_back (static_cast <double >(fitInfo.ampFT0A ));
617634 valuesToFill.push_back (static_cast <double >(fitInfo.ampFT0C ));
@@ -630,6 +647,12 @@ struct HfTaskD0 {
630647 auto fillTHnData = [&](float mass, int d0Type) {
631648 std::vector<double > valuesToFill{static_cast <double >(mass), static_cast <double >(ptCandidate),
632649 static_cast <double >(HfHelper::yD0 (candidate)), static_cast <double >(d0Type)};
650+ if (storeCentrality) {
651+ valuesToFill.push_back (cent);
652+ }
653+ if (storeOccupancyAndIR) {
654+ valuesToFill.push_back (occ);
655+ }
633656 valuesToFill.push_back (static_cast <double >(gapTypeInt));
634657 valuesToFill.push_back (static_cast <double >(fitInfo.ampFT0A ));
635658 valuesToFill.push_back (static_cast <double >(fitInfo.ampFT0C ));
0 commit comments