@@ -335,20 +335,29 @@ struct HfTaskD0 {
335335 axes.push_back (thnAxisMinItsNCls);
336336 axes.push_back (thnAxisMinTpcNCrossedRows);
337337 }
338- if (doprocessDataWithDCAFitterNMlWithUpc) {
339- axes.push_back (thnAxisGapType);
340- axes.push_back (thnAxisFT0A);
341- axes.push_back (thnAxisFT0C);
342- }
343338 if (applyMl) {
344339 const AxisSpec thnAxisBkgScore{thnConfigAxisBkgScore, " BDT score bkg." };
345340 const AxisSpec thnAxisNonPromptScore{thnConfigAxisNonPromptScore, " BDT score non-prompt." };
346341 const AxisSpec thnAxisPromptScore{thnConfigAxisPromptScore, " BDT score prompt." };
347342
348- axes.insert (axes.begin (), thnAxisPromptScore);
349- axes.insert (axes.begin (), thnAxisNonPromptScore);
350- axes.insert (axes.begin (), thnAxisBkgScore);
343+ // Insert ML scores after pt (position 2) to match taskDplus structure: [mass, pt, mlScores, ...]
344+ if (doprocessDataWithDCAFitterNMlWithUpc) {
345+ axes.insert (axes.begin () + 2 , thnAxisPromptScore);
346+ axes.insert (axes.begin () + 2 , thnAxisNonPromptScore);
347+ axes.insert (axes.begin () + 2 , thnAxisBkgScore);
348+ } else {
349+ axes.insert (axes.begin (), thnAxisPromptScore);
350+ axes.insert (axes.begin (), thnAxisNonPromptScore);
351+ axes.insert (axes.begin (), thnAxisBkgScore);
352+ }
353+ }
354+ if (doprocessDataWithDCAFitterNMlWithUpc) {
355+ axes.push_back (thnAxisGapType);
356+ axes.push_back (thnAxisFT0A);
357+ axes.push_back (thnAxisFT0C);
358+ }
351359
360+ if (applyMl) {
352361 registry.add (" hBdtScoreVsMassVsPtVsPtBVsYVsOriginVsD0Type" , " Thn for D0 candidates" , HistType::kTHnSparseD , axes);
353362 registry.get <THnSparse>(HIST (" hBdtScoreVsMassVsPtVsPtBVsYVsOriginVsD0Type" ))->Sumw2 ();
354363 } else {
@@ -581,6 +590,8 @@ struct HfTaskD0 {
581590 int const gapTypeInt = hf_upc::gapTypeToInt (gap);
582591 const auto thisCollId = collision.globalIndex ();
583592 const auto & groupedD0Candidates = candidates.sliceBy (candD0PerCollision, thisCollId);
593+ float cent{centrality};
594+ float occ{-1 .f };
584595
585596 for (const auto & candidate : groupedD0Candidates) {
586597 if (!(candidate.hfflag () & 1 << aod::hf_cand_2prong::DecayType::D0ToPiK)) {
@@ -605,12 +616,18 @@ struct HfTaskD0 {
605616 registry.fill (HIST (" hMassVsPhi" ), massD0bar, ptCandidate, candidate.phi ());
606617 }
607618
608- // Fill THnSparse with gap type and FIT signals using vectorized approach
619+ // Fill THnSparse with structure matching taskDplus: [mass, pt, mlScores, cent, occ, gapType, FT0A, FT0C]
609620 if constexpr (fillMl) {
610621 auto fillTHnData = [&](float mass, int d0Type) {
611- std::vector<double > valuesToFill{candidate. mlProbD0 ()[ 0 ], candidate. mlProbD0 ()[ 1 ], candidate. mlProbD0 ()[ 2 ] ,
612- static_cast < double >(mass), static_cast < double >(ptCandidate) ,
622+ std::vector<double > valuesToFill{static_cast < double >(mass), static_cast < double >(ptCandidate) ,
623+ candidate. mlProbD0 ()[ 0 ], candidate. mlProbD0 ()[ 1 ], candidate. mlProbD0 ()[ 2 ] ,
613624 static_cast <double >(hfHelper.yD0 (candidate)), static_cast <double >(d0Type)};
625+ if (storeCentrality) {
626+ valuesToFill.push_back (cent);
627+ }
628+ if (storeOccupancyAndIR) {
629+ valuesToFill.push_back (occ);
630+ }
614631 valuesToFill.push_back (static_cast <double >(gapTypeInt));
615632 valuesToFill.push_back (static_cast <double >(fitInfo.ampFT0A ));
616633 valuesToFill.push_back (static_cast <double >(fitInfo.ampFT0C ));
@@ -629,6 +646,12 @@ struct HfTaskD0 {
629646 auto fillTHnData = [&](float mass, int d0Type) {
630647 std::vector<double > valuesToFill{static_cast <double >(mass), static_cast <double >(ptCandidate),
631648 static_cast <double >(hfHelper.yD0 (candidate)), static_cast <double >(d0Type)};
649+ if (storeCentrality) {
650+ valuesToFill.push_back (cent);
651+ }
652+ if (storeOccupancyAndIR) {
653+ valuesToFill.push_back (occ);
654+ }
632655 valuesToFill.push_back (static_cast <double >(gapTypeInt));
633656 valuesToFill.push_back (static_cast <double >(fitInfo.ampFT0A ));
634657 valuesToFill.push_back (static_cast <double >(fitInfo.ampFT0C ));
0 commit comments