Skip to content

Commit 8789b18

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 750f63e commit 8789b18

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
@@ -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

Comments
 (0)