Skip to content

Commit 1c10742

Browse files
minjungkim12claude
andcommitted
[PWGHF] Unify taskD0 THnSparse structure with taskDplus
- Reorder axes to match taskDplus: [mass, pt, mlScores, y, d0Type, cent, occ, gapType, FT0A, FT0C] - Insert ML scores after pt (position 2) instead of at the beginning - Add centrality and occupancy support in UPC filling function - Conditional filling based on storeCentrality and storeOccupancyAndIR flags - Consistent THnSparse structure across both D0 and Dplus tasks - Enables flexible analysis with optional centrality/occupancy axes 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
1 parent b91499a commit 1c10742

File tree

1 file changed

+34
-11
lines changed

1 file changed

+34
-11
lines changed

PWGHF/D2H/Tasks/taskD0.cxx

Lines changed: 34 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -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

Comments
 (0)