Skip to content

Commit a63fb5c

Browse files
minjungkim12claude
andcommitted
[PWGHF] Refactor UPC THnSparse filling: use vectorized approach with lambda
- Refactor both taskD0 and taskDplus to use consistent vector-based filling - Add lambda function fillTHnData for cleaner, more maintainable code - Build valuesToFill vector dynamically and push gap type - Revert fillSparseML modifications in taskDplus (remove gapType param) - Replace explicit registry.fill() calls with THnSparse->Fill(vector) - Consistent with taskLc implementation pattern - Reduces code duplication and improves readability 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude <noreply@anthropic.com>
1 parent 04260f6 commit a63fb5c

File tree

2 files changed

+49
-37
lines changed

2 files changed

+49
-37
lines changed

PWGHF/D2H/Tasks/taskD0.cxx

Lines changed: 24 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -599,23 +599,39 @@ struct HfTaskD0 {
599599
registry.fill(HIST("hMassVsPhi"), massD0bar, ptCandidate, candidate.phi());
600600
}
601601

602+
// Fill THnSparse with gap type using vectorized approach similar to taskDplus
602603
if constexpr (fillMl) {
604+
auto fillTHnData = [&](float mass, int d0Type) {
605+
std::vector<double> valuesToFill{candidate.mlProbD0()[0], candidate.mlProbD0()[1], candidate.mlProbD0()[2],
606+
static_cast<double>(mass), static_cast<double>(ptCandidate),
607+
static_cast<double>(hfHelper.yD0(candidate)), static_cast<double>(d0Type)};
608+
valuesToFill.push_back(static_cast<double>(gapTypeInt));
609+
registry.get<THnSparse>(HIST("hBdtScoreVsMassVsPtVsPtBVsYVsOriginVsD0Type"))->Fill(valuesToFill.data());
610+
};
611+
603612
if (candidate.isSelD0() >= selectionFlagD0) {
604-
registry.fill(HIST("hBdtScoreVsMassVsPtVsPtBVsYVsOriginVsD0Type"), candidate.mlProbD0()[0], candidate.mlProbD0()[1], candidate.mlProbD0()[2], massD0, ptCandidate, hfHelper.yD0(candidate), SigD0, static_cast<float>(gapTypeInt));
605-
registry.fill(HIST("hBdtScoreVsMassVsPtVsPtBVsYVsOriginVsD0Type"), candidate.mlProbD0()[0], candidate.mlProbD0()[1], candidate.mlProbD0()[2], massD0, ptCandidate, hfHelper.yD0(candidate), candidate.isSelD0bar() ? ReflectedD0 : PureSigD0, static_cast<float>(gapTypeInt));
613+
fillTHnData(massD0, SigD0);
614+
fillTHnData(massD0, candidate.isSelD0bar() ? ReflectedD0 : PureSigD0);
606615
}
607616
if (candidate.isSelD0bar() >= selectionFlagD0bar) {
608-
registry.fill(HIST("hBdtScoreVsMassVsPtVsPtBVsYVsOriginVsD0Type"), candidate.mlProbD0()[0], candidate.mlProbD0()[1], candidate.mlProbD0()[2], massD0bar, ptCandidate, hfHelper.yD0(candidate), SigD0bar, static_cast<float>(gapTypeInt));
609-
registry.fill(HIST("hBdtScoreVsMassVsPtVsPtBVsYVsOriginVsD0Type"), candidate.mlProbD0()[0], candidate.mlProbD0()[1], candidate.mlProbD0()[2], massD0bar, ptCandidate, hfHelper.yD0(candidate), candidate.isSelD0() ? ReflectedD0bar : PureSigD0bar, static_cast<float>(gapTypeInt));
617+
fillTHnData(massD0bar, SigD0bar);
618+
fillTHnData(massD0bar, candidate.isSelD0() ? ReflectedD0bar : PureSigD0bar);
610619
}
611620
} else {
621+
auto fillTHnData = [&](float mass, int d0Type) {
622+
std::vector<double> valuesToFill{static_cast<double>(mass), static_cast<double>(ptCandidate),
623+
static_cast<double>(hfHelper.yD0(candidate)), static_cast<double>(d0Type)};
624+
valuesToFill.push_back(static_cast<double>(gapTypeInt));
625+
registry.get<THnSparse>(HIST("hMassVsPtVsPtBVsYVsOriginVsD0Type"))->Fill(valuesToFill.data());
626+
};
627+
612628
if (candidate.isSelD0() >= selectionFlagD0) {
613-
registry.fill(HIST("hMassVsPtVsPtBVsYVsOriginVsD0Type"), massD0, ptCandidate, hfHelper.yD0(candidate), SigD0, static_cast<float>(gapTypeInt));
614-
registry.fill(HIST("hMassVsPtVsPtBVsYVsOriginVsD0Type"), massD0, ptCandidate, hfHelper.yD0(candidate), candidate.isSelD0bar() ? ReflectedD0 : PureSigD0, static_cast<float>(gapTypeInt));
629+
fillTHnData(massD0, SigD0);
630+
fillTHnData(massD0, candidate.isSelD0bar() ? ReflectedD0 : PureSigD0);
615631
}
616632
if (candidate.isSelD0bar() >= selectionFlagD0bar) {
617-
registry.fill(HIST("hMassVsPtVsPtBVsYVsOriginVsD0Type"), massD0bar, ptCandidate, hfHelper.yD0(candidate), SigD0bar, static_cast<float>(gapTypeInt));
618-
registry.fill(HIST("hMassVsPtVsPtBVsYVsOriginVsD0Type"), massD0bar, ptCandidate, hfHelper.yD0(candidate), candidate.isSelD0() ? ReflectedD0bar : PureSigD0bar, static_cast<float>(gapTypeInt));
633+
fillTHnData(massD0bar, SigD0bar);
634+
fillTHnData(massD0bar, candidate.isSelD0() ? ReflectedD0bar : PureSigD0bar);
619635
}
620636
}
621637
}

PWGHF/D2H/Tasks/taskDplus.cxx

Lines changed: 25 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -314,15 +314,13 @@ struct HfTaskDplus {
314314
/// \param centrality collision centrality
315315
/// \param occupancy collision occupancy
316316
/// \param numPvContributors contributors to the PV
317-
/// \param gapType UPC gap type (-1 for non-UPC)
318317
template <bool IsMc, bool IsMatched, typename T1>
319318
void fillSparseML(const T1& candidate,
320319
float ptbhad,
321320
int flagBHad,
322321
float centrality,
323322
float occupancy,
324-
float numPvContributors,
325-
int gapType = -1)
323+
float numPvContributors)
326324
{
327325
std::vector<float> outputMl = {-999., -999., -999.};
328326
for (unsigned int iclass = 0; iclass < classMl->size(); iclass++) {
@@ -387,30 +385,16 @@ struct HfTaskDplus {
387385
}
388386
}
389387
} else { // Data
390-
if (gapType >= 0) {
391-
// UPC mode: always include gap type
392-
if (storeCentrality && storeOccupancy) {
393-
registry.fill(HIST("hSparseMass"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2], centrality, occupancy, static_cast<float>(gapType));
394-
} else if (storeCentrality && !storeOccupancy) {
395-
registry.fill(HIST("hSparseMass"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2], centrality, static_cast<float>(gapType));
396-
} else if (!storeCentrality && storeOccupancy) {
397-
registry.fill(HIST("hSparseMass"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2], occupancy, static_cast<float>(gapType));
398-
} else {
399-
registry.fill(HIST("hSparseMass"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2], static_cast<float>(gapType));
400-
}
388+
if (storeCentrality && storeOccupancy) {
389+
registry.fill(HIST("hSparseMass"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2], centrality, occupancy);
390+
} else if (storeCentrality && !storeOccupancy) {
391+
registry.fill(HIST("hSparseMass"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2], centrality);
392+
} else if (!storeCentrality && storeOccupancy) {
393+
registry.fill(HIST("hSparseMass"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2], occupancy);
394+
} else if (!storeCentrality && !storeOccupancy && storePvContributors) {
395+
registry.fill(HIST("hSparseMass"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2], numPvContributors);
401396
} else {
402-
// Non-UPC mode: original behavior
403-
if (storeCentrality && storeOccupancy) {
404-
registry.fill(HIST("hSparseMass"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2], centrality, occupancy);
405-
} else if (storeCentrality && !storeOccupancy) {
406-
registry.fill(HIST("hSparseMass"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2], centrality);
407-
} else if (!storeCentrality && storeOccupancy) {
408-
registry.fill(HIST("hSparseMass"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2], occupancy);
409-
} else if (!storeCentrality && !storeOccupancy && storePvContributors) {
410-
registry.fill(HIST("hSparseMass"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2], numPvContributors);
411-
} else {
412-
registry.fill(HIST("hSparseMass"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2]);
413-
}
397+
registry.fill(HIST("hSparseMass"), hfHelper.invMassDplusToPiKPi(candidate), candidate.pt(), outputMl[0], outputMl[1], outputMl[2]);
414398
}
415399
}
416400
}
@@ -732,8 +716,6 @@ struct HfTaskDplus {
732716
float cent{-1.f};
733717
float occ{-1.f};
734718
float numPvContr{-1.f};
735-
float ptBhad{-1.f};
736-
int const flagBHad{-1};
737719
int const gapTypeInt = hf_upc::gapTypeToInt(gap);
738720

739721
for (const auto& candidate : groupedDplusCandidates) {
@@ -742,7 +724,21 @@ struct HfTaskDplus {
742724
}
743725
fillHisto(candidate);
744726
if constexpr (fillMl) {
745-
fillSparseML<false, false>(candidate, ptBhad, flagBHad, cent, occ, numPvContr, gapTypeInt);
727+
// Fill THn with gap type using lambda function similar to taskLc
728+
std::vector<float> outputMl = {-999., -999., -999.};
729+
for (unsigned int iclass = 0; iclass < classMl->size(); iclass++) {
730+
outputMl[iclass] = candidate.mlProbDplusToPiKPi()[classMl->at(iclass)];
731+
}
732+
733+
std::vector<double> valuesToFill{hfHelper.invMassDplusToPiKPi(candidate), static_cast<double>(candidate.pt()), outputMl[0], outputMl[1], outputMl[2]};
734+
if (storeCentrality) {
735+
valuesToFill.push_back(cent);
736+
}
737+
if (storeOccupancy) {
738+
valuesToFill.push_back(occ);
739+
}
740+
valuesToFill.push_back(static_cast<double>(gapTypeInt));
741+
registry.get<THnSparse>(HIST("hSparseMass"))->Fill(valuesToFill.data());
746742
}
747743
}
748744
}

0 commit comments

Comments
 (0)