Skip to content

Commit e744805

Browse files
[PWGHF/D2H] Adding ML application in MC processes for analysis on Multiplicity Dependent D* Production Cross Section (AliceO2Group#9119)
1 parent d7f46ce commit e744805

File tree

1 file changed

+109
-59
lines changed

1 file changed

+109
-59
lines changed

PWGHF/D2H/Tasks/taskDstarToD0Pi.cxx

Lines changed: 109 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@
1515
/// \author Deependra Sharma <[email protected]>, IITB
1616
/// \author Fabrizio Grosa <[email protected]>, CERN
1717

18+
/// \brief Dstar production analysis task (With and Without ML)
19+
1820
#include <algorithm>
1921
#include <utility>
2022
#include <vector>
@@ -42,11 +44,14 @@ struct HfTaskDstarToD0Pi {
4244
Configurable<bool> selectionFlagHfD0ToPiK{"selectionFlagHfD0ToPiK", true, "Selection Flag for HF flagged candidates"};
4345
Configurable<std::vector<double>> ptBins{"ptBins", std::vector<double>{hf_cuts_dstar_to_d0_pi::vecBinsPt}, "pT bin limits for Dstar"};
4446

47+
SliceCache cache;
48+
4549
using CandDstarWSelFlag = soa::Join<aod::HfD0FromDstar, aod::HfCandDstars, aod::HfSelDstarToD0Pi>;
4650
using CandDstarWSelFlagWMl = soa::Join<aod::HfD0FromDstar, aod::HfCandDstars, aod::HfSelDstarToD0Pi, aod::HfMlDstarToD0Pi>;
4751
/// @brief specially for MC data
4852
// full reconstructed Dstar candidate
4953
using CandDstarWSelFlagMcRec = soa::Join<aod::HfD0FromDstar, aod::HfCandDstars, aod::HfSelDstarToD0Pi, aod::HfCandDstarMcRec>;
54+
using CandDstarWSelFlagWMlMcRec = soa::Join<aod::HfD0FromDstar, aod::HfCandDstars, aod::HfSelDstarToD0Pi, aod::HfCandDstarMcRec, aod::HfMlDstarToD0Pi>;
5055
using CandDstarMcGen = soa::Join<aod::McParticles, aod::HfCandDstarMcGen>;
5156

5257
using CollisionsWCent = soa::Join<aod::Collisions, aod::CentFT0Ms>;
@@ -57,9 +62,9 @@ struct HfTaskDstarToD0Pi {
5762
Preslice<soa::Filtered<CandDstarWSelFlagWMl>> preslicSelectedCandDstarPerColWMl = aod::hf_cand::collisionId;
5863

5964
PresliceUnsorted<CollisionsWCentMcLabel> colsPerMcCollision = aod::mccollisionlabel::mcCollisionId;
60-
SliceCache cache;
6165

6266
Partition<CandDstarWSelFlagMcRec> rowsSelectedCandDstarMcRec = aod::hf_sel_candidate_dstar::isRecoD0Flag == selectionFlagHfD0ToPiK;
67+
Partition<CandDstarWSelFlagWMlMcRec> rowsSelectedCandDstarMcRecWMl = aod::hf_sel_candidate_dstar::isRecoD0Flag == selectionFlagHfD0ToPiK;
6368

6469
ConfigurableAxis binningImpactParam{"binningImpactParam", {1000, 0.1, -0.1}, " Bins of Impact Parameter"};
6570
ConfigurableAxis binningDecayLength{"binningDecayLength", {1000, 0.0, 0.7}, "Bins of Decay Length"};
@@ -82,6 +87,9 @@ struct HfTaskDstarToD0Pi {
8287

8388
void init(InitContext&)
8489
{
90+
if ((doprocessDataWoML && doprocessDataWML) || (doprocessMcWoMl && doprocessMcWML) || (doprocessDataWoML && doprocessMcWML) || (doprocessDataWML && doprocessMcWoMl)) {
91+
LOGP(fatal, "Only Without-ML or With-ML functions should be enabled at a time! Please check your configuration!");
92+
}
8593
auto vecPtBins = (std::vector<double>)ptBins;
8694

8795
AxisSpec axisImpactParam = {binningImpactParam, "impact parameter (cm)"};
@@ -162,9 +170,29 @@ struct HfTaskDstarToD0Pi {
162170
registry.add("Efficiency/hNumPvContributorsCandInMass", "PV Contributors; PV Contributor; FT0M Centrality", {HistType::kTH2F, {{100, 0, 300}, {axisCentrality}}}, true);
163171

164172
// BDT Score (axisBDTScoreBackground, axisBDTScorePrompt, axisBDTScoreNonPrompt)
165-
registry.add("Yield/hDeltaInvMassVsPtVsCentVsBDTScore", "#Delta #it{M}_{inv} Vs Pt Vs Cent Vs BDTScore", {HistType::kTHnSparseL, {{axisDeltaInvMass}, {vecPtBins, "#it{p}_{T} (GeV/#it{c})"}, {axisCentrality}, {axisBDTScoreBackground}, {axisBDTScorePrompt}, {axisBDTScoreNonPrompt}}});
173+
if (doprocessDataWML) {
174+
registry.add("Yield/hDeltaInvMassVsPtVsCentVsBDTScore", "#Delta #it{M}_{inv} Vs Pt Vs Cent Vs BDTScore", {HistType::kTHnSparseF, {{axisDeltaInvMass}, {vecPtBins, "#it{p}_{T} (GeV/#it{c})"}, {axisCentrality}, {axisBDTScoreBackground}, {axisBDTScorePrompt}, {axisBDTScoreNonPrompt}}});
175+
}
176+
if (doprocessMcWML) {
177+
registry.add("Efficiency/hPtPromptVsCentVsBDTScore", "Pt Vs Cent Vs BDTScore", {HistType::kTHnSparseF, {{vecPtBins, "#it{p}_{T} (GeV/#it{c})"}, {axisCentrality}, {axisBDTScoreBackground}, {axisBDTScorePrompt}, {axisBDTScoreNonPrompt}}});
178+
registry.add("Efficiency/hPtNonPromptVsCentVsBDTScore", "Pt Vs Cent Vs BDTScore", {HistType::kTHnSparseF, {{vecPtBins, "#it{p}_{T} (GeV/#it{c})"}, {axisCentrality}, {axisBDTScoreBackground}, {axisBDTScorePrompt}, {axisBDTScoreNonPrompt}}});
179+
// registry.add("Efficiency/hPtBkgVsCentVsBDTScore", "Pt Vs Cent Vs BDTScore", {HistType::kTHnSparseF, {{vecPtBins, "#it{p}_{T} (GeV/#it{c})"}, {axisCentrality}, {axisBDTScoreBackground}, {axisBDTScorePrompt}, {axisBDTScoreNonPrompt}}});
180+
}
166181
}
167182

183+
// Comparator function to sort based on the second argument of a tuple
184+
static bool compare(const std::pair<soa::Filtered<CollisionsWCentMcLabel>::iterator, int>& a, const std::pair<soa::Filtered<CollisionsWCentMcLabel>::iterator, int>& b)
185+
{
186+
return a.second > b.second;
187+
}
188+
189+
/// @brief This function runs over Data to obatin yield
190+
/// @tparam T1 type of the candidate
191+
/// @tparam T2 type of preslice used to slice the candidate table
192+
/// @tparam applyMl a boolean to apply ML or not
193+
/// @param cols reconstructed collision with centrality
194+
/// @param selectedCands selected candidates with selection flag
195+
/// @param preslice preslice to slice
168196
template <bool applyMl, typename T1, typename T2>
169197
void runTaskDstar(CollisionsWCent const& cols, T1 selectedCands, T2 preslice)
170198
{
@@ -178,7 +206,7 @@ struct HfTaskDstarToD0Pi {
178206
auto nCandsCurrentCol = selectedCandsCurrentCol.size();
179207

180208
if (nCandsCurrentCol > 0) {
181-
LOGF(debug, "size of selectedCandsCurrentCol: %d", nCandsCurrentCol);
209+
// LOGF(debug, "size of selectedCandsCurrentCol: %d", nCandsCurrentCol);
182210
registry.fill(HIST("Efficiency/hNumPvContributorsCand"), nPVContributors, centrality);
183211
}
184212

@@ -259,47 +287,34 @@ struct HfTaskDstarToD0Pi {
259287
} // collision loop ends
260288
}
261289

262-
// process function without susing ML
263-
void processWoML(CollisionsWCent const& cols, soa::Filtered<CandDstarWSelFlag> const& selectedCands)
290+
/// @brief This function runs over MC at reco level to obatin efficiency
291+
/// @tparam T1 type of the candidate table
292+
/// @tparam applyMl a boolean to apply ML or not
293+
/// @param candsMcRecSel reconstructed candidates with selection flag
294+
/// @param rowsMcPartilces generated particles table
295+
template <bool applyMl, typename T1>
296+
void runMcRecTaskDstar(T1 const& candsMcRecSel, CandDstarMcGen const& rowsMcPartilces)
264297
{
265-
runTaskDstar<false, soa::Filtered<CandDstarWSelFlag>, Preslice<soa::Filtered<CandDstarWSelFlag>>>(cols, selectedCands, preslicSelectedCandDstarPerCol);
266-
}
267-
PROCESS_SWITCH(HfTaskDstarToD0Pi, processWoML, "Process without ML", true);
268-
269-
// process function with susing ML, Here we store BDT score as well
270-
void processWML(CollisionsWCent const& cols, soa::Filtered<CandDstarWSelFlagWMl> const& selectedCands)
271-
{
272-
runTaskDstar<true, soa::Filtered<CandDstarWSelFlagWMl>, Preslice<soa::Filtered<CandDstarWSelFlagWMl>>>(cols, selectedCands, preslicSelectedCandDstarPerColWMl);
273-
}
274-
PROCESS_SWITCH(HfTaskDstarToD0Pi, processWML, "Process with ML", false);
275-
276-
void processMC(aod::McCollisions const&, CollisionsWCentMcLabel const& collisions, CandDstarWSelFlagMcRec const&,
277-
CandDstarMcGen const& rowsMcPartilces,
278-
aod::TracksWMc const&)
279-
{
280-
rowsSelectedCandDstarMcRec.bindExternalIndices(&collisions);
281-
298+
// LOGF(info, "Running MC Rec Task Dstar");
282299
int8_t signDstar = 0;
283300
// MC at Reconstruction level
284-
for (const auto& candDstarMcRec : rowsSelectedCandDstarMcRec) {
301+
for (const auto& candDstarMcRec : candsMcRecSel) {
302+
// LOGF(info, "MC Rec Dstar loop");
285303
auto ptDstarRecSig = candDstarMcRec.pt();
286304
auto yDstarRecSig = candDstarMcRec.y(constants::physics::MassDStar);
287305
if (yCandDstarRecoMax >= 0. && std::abs(yDstarRecSig) > yCandDstarRecoMax) {
288306
continue;
289307
}
290-
auto collision = candDstarMcRec.collision_as<CollisionsWCentMcLabel>();
308+
auto collision = candDstarMcRec.template collision_as<CollisionsWCentMcLabel>();
291309
auto centrality = collision.centFT0M(); // 0-100%
292310
if (TESTBIT(std::abs(candDstarMcRec.flagMcMatchRec()), aod::hf_cand_dstar::DecayType::DstarToD0Pi)) { // if MC matching is successful at Reconstruction Level
311+
// LOGF(info, "MC Rec Dstar loop MC Matched");
293312
// get MC Mother particle
294-
auto indexMother = RecoDecay::getMother(rowsMcPartilces, candDstarMcRec.prong0_as<aod::TracksWMc>().mcParticle_as<CandDstarMcGen>(), o2::constants::physics::Pdg::kDStar, true, &signDstar, 2);
313+
auto prong0 = candDstarMcRec.template prong0_as<aod::TracksWMc>();
314+
auto indexMother = RecoDecay::getMother(rowsMcPartilces, prong0.template mcParticle_as<CandDstarMcGen>(), o2::constants::physics::Pdg::kDStar, true, &signDstar, 2);
295315
auto particleMother = rowsMcPartilces.rawIteratorAt(indexMother); // What is difference between rawIterator() or iteratorAt() methods?
296316
registry.fill(HIST("QA/hPtSkimDstarGenSig"), particleMother.pt()); // generator level pt
297317
registry.fill(HIST("Efficiency/hPtVsCentSkimDstarGenSig"), particleMother.pt(), centrality);
298-
registry.fill(HIST("Efficiency/hPtVsCentFullRecoDstarRecSig"), ptDstarRecSig, centrality);
299-
300-
// auto recCollision = candDstarMcRec.collision_as<CollisionsWCentMcLabel>();
301-
// float centFT0M = recCollision.centFT0M();
302-
// LOGF(info, "centFT0M: %f", centFT0M);
303318

304319
registry.fill(HIST("QA/hPtVsYSkimDstarRecSig"), ptDstarRecSig, yDstarRecSig); // Skimed at level of trackIndexSkimCreator
305320
if (candDstarMcRec.isRecoTopol()) { // if Topological selection are passed
@@ -308,8 +323,9 @@ struct HfTaskDstarToD0Pi {
308323
if (candDstarMcRec.isRecoPid()) { // if PID selection is passed
309324
registry.fill(HIST("QA/hPtVsYRecoPidDstarRecSig"), ptDstarRecSig, yDstarRecSig);
310325
}
311-
if (candDstarMcRec.isRecoCand()) { // if all selection passed
326+
if (candDstarMcRec.isSelDstarToD0Pi()) { // if all selection passed
312327
registry.fill(HIST("QA/hPtFullRecoDstarRecSig"), ptDstarRecSig);
328+
registry.fill(HIST("Efficiency/hPtVsCentFullRecoDstarRecSig"), ptDstarRecSig, centrality);
313329
}
314330
registry.fill(HIST("QA/hCPASkimD0RecSig"), candDstarMcRec.cpaD0());
315331
registry.fill(HIST("QA/hEtaSkimD0RecSig"), candDstarMcRec.etaD0());
@@ -324,76 +340,81 @@ struct HfTaskDstarToD0Pi {
324340
if (candDstarMcRec.isRecoPid()) { // if PID selection is passed
325341
registry.fill(HIST("QA/hPtVsYRecoPidPromptDstarRecSig"), ptDstarRecSig, yDstarRecSig);
326342
}
327-
if (candDstarMcRec.isRecoCand()) { // if all selection passed
343+
if (candDstarMcRec.isSelDstarToD0Pi()) { // if all selection passed
328344
registry.fill(HIST("QA/hPtFullRecoPromptDstarRecSig"), ptDstarRecSig);
345+
if constexpr (applyMl) {
346+
// LOGF(info, "Deep: Prompt MC Rec Task Dstar: ML applied");
347+
auto bdtScore = candDstarMcRec.mlProbDstarToD0Pi();
348+
registry.fill(HIST("Efficiency/hPtPromptVsCentVsBDTScore"), ptDstarRecSig, centrality, bdtScore[0], bdtScore[1], bdtScore[2]);
349+
// LOGF(info, "Deep: Prompt MC Rec Task Dstar: ML applied, Efficiency filled");
350+
}
329351
}
330-
331352
} else if (candDstarMcRec.originMcRec() == RecoDecay::OriginType::NonPrompt) { // only non-prompt signal at reconstruction level
332353
registry.fill(HIST("QA/hPtVsYSkimNonPromptDstarRecSig"), ptDstarRecSig, yDstarRecSig);
333354
if (candDstarMcRec.isRecoTopol()) { // if Topological selection are passed
334355
registry.fill(HIST("QA/hPtVsYRecoTopolNonPromptDstarRecSig"), ptDstarRecSig, yDstarRecSig);
335356
}
336-
if (candDstarMcRec.isRecoPid()) {
357+
if (candDstarMcRec.isRecoPid()) { // if PID selection is passed
337358
registry.fill(HIST("QA/hPtVsYRecoPidNonPromptDstarRecSig"), ptDstarRecSig, yDstarRecSig);
338359
}
339-
if (candDstarMcRec.isRecoCand()) {
360+
if (candDstarMcRec.isSelDstarToD0Pi()) { // if all selection passed
340361
registry.fill(HIST("QA/hPtFullRecoNonPromptDstarRecSig"), ptDstarRecSig);
362+
if constexpr (applyMl) {
363+
auto bdtScore = candDstarMcRec.mlProbDstarToD0Pi();
364+
registry.fill(HIST("Efficiency/hPtNonPromptVsCentVsBDTScore"), ptDstarRecSig, centrality, bdtScore[0], bdtScore[1], bdtScore[2]);
365+
}
341366
}
342367
}
343368
} else { // MC Unmatched (Baground at Reconstruction Level)
344369
registry.fill(HIST("QA/hCPASkimD0RecBg"), candDstarMcRec.cpaD0());
345370
registry.fill(HIST("QA/hEtaSkimD0RecBg"), candDstarMcRec.etaD0());
346371
registry.fill(HIST("QA/hEtaSkimDstarRecBg"), candDstarMcRec.eta());
347-
registry.fill(HIST("QA/hPtSkimD0RecBg"), candDstarMcRec.ptD0());
348-
registry.fill(HIST("QA/hPtSkimDstarRecBg"), candDstarMcRec.pt());
372+
if (candDstarMcRec.isSelDstarToD0Pi()) {
373+
registry.fill(HIST("QA/hPtSkimDstarRecBg"), candDstarMcRec.pt());
374+
}
349375
}
350-
}
376+
} // candidate loop ends
377+
// LOGF(info, "Deep: MC Rec Task Dstar finished");
378+
}
351379

352-
// MC at Generator Level
353-
for (const auto& mcParticle : rowsMcPartilces) {
380+
/// @brief This function runs over MC at gen level to obatin efficiency
381+
/// @param collisions reconstructed collision with centrality
382+
/// @param rowsMcPartilces generated particles table
383+
void runMcGenTaskDstar(CollisionsWCentMcLabel const& collisions, CandDstarMcGen const& rowsMcPartilces)
384+
{
385+
// MC Gen level
386+
for (auto const& mcParticle : rowsMcPartilces) {
354387
if (TESTBIT(std::abs(mcParticle.flagMcMatchGen()), aod::hf_cand_dstar::DecayType::DstarToD0Pi)) { // MC Matching is successful at Generator Level
355-
356388
auto ptGen = mcParticle.pt();
357-
// auto yGen = mcParticle.y(); // Can we use this definition?
358389
auto yGen = RecoDecay::y(mcParticle.pVector(), o2::constants::physics::MassDStar);
359390
if (yCandDstarGenMax >= 0. && std::abs(yGen) > yCandDstarGenMax) {
360391
continue;
361392
}
362-
363-
auto mcCollision = mcParticle.mcCollision_as<aod::McCollisions>();
393+
auto mcCollision = mcParticle.template mcCollision_as<aod::McCollisions>();
364394
auto recCollisions = collisions.sliceBy(colsPerMcCollision, mcCollision.globalIndex());
365-
// auto recCollisions = collisions.sliceByCached(aod::mccollisionlabel::mcCollisionId, mcCollision.globalIndex(), cache);
366-
// auto recCollisions = collisions.sliceByCachedUnsorted(aod::mccollisionlabel::mcCollisionId, mcCollision.globalIndex(), cache);
367-
368395
// looking if a generated collision reconstructed more than a times.
369396
if (recCollisions.size() > 1) {
370397
for (const auto& [c1, c2] : combinations(CombinationsStrictlyUpperIndexPolicy(recCollisions, recCollisions))) {
371-
auto deltaCent = abs(c1.centFT0M() - c2.centFT0M());
398+
auto deltaCent = std::abs(c1.centFT0M() - c2.centFT0M());
372399
registry.fill(HIST("QA/hDeltaCentGen"), deltaCent);
373400
}
374401
}
375-
376402
float centFT0MGen;
377403
// assigning centrality to MC Collision using max FT0M amplitute from Reconstructed collisions
378404
if (recCollisions.size()) {
379405
std::vector<std::pair<soa::Filtered<CollisionsWCentMcLabel>::iterator, int>> tempRecCols;
380-
381406
for (const auto& recCol : recCollisions) {
382-
// if(recCollisions.size()>1) LOGF(info, "cuurent cent: %f",recCol.centFT0M());
383407
tempRecCols.push_back(std::make_pair(recCol, recCol.numContrib()));
384408
}
385409
std::sort(tempRecCols.begin(), tempRecCols.end(), compare);
386410
centFT0MGen = tempRecCols.at(0).first.centFT0M();
387-
// if(recCollisions.size()>1) LOGF(info, "assigned cent: %f",centFT0MGen);
388411
} else {
389412
centFT0MGen = -999.;
390413
}
391-
392414
registry.fill(HIST("QA/hEtaDstarGen"), mcParticle.eta());
393415
registry.fill(HIST("QA/hPtDstarGen"), ptGen);
394416
registry.fill(HIST("QA/hPtVsYDstarGen"), ptGen, yGen);
395417
registry.fill(HIST("Efficiency/hPtVsCentDstarGen"), ptGen, centFT0MGen);
396-
397418
// only promt Dstar candidate at Generator level
398419
if (mcParticle.originMcGen() == RecoDecay::OriginType::Prompt) {
399420
registry.fill(HIST("QA/hPtPromptDstarGen"), ptGen);
@@ -403,15 +424,44 @@ struct HfTaskDstarToD0Pi {
403424
registry.fill(HIST("QA/hPtVsYNonPromptDstarGen"), ptGen, yGen);
404425
}
405426
}
406-
}
427+
} // MC Particle loop ends
407428
}
408-
PROCESS_SWITCH(HfTaskDstarToD0Pi, processMC, "Process MC Data", false);
409429

410-
// Comparator function to sort based on the second argument of a tuple
411-
static bool compare(const std::pair<soa::Filtered<CollisionsWCentMcLabel>::iterator, int>& a, const std::pair<soa::Filtered<CollisionsWCentMcLabel>::iterator, int>& b)
430+
// process data function without susing ML
431+
void processDataWoML(CollisionsWCent const& cols, soa::Filtered<CandDstarWSelFlag> const& selectedCands)
412432
{
413-
return a.second > b.second;
433+
runTaskDstar<false, soa::Filtered<CandDstarWSelFlag>, Preslice<soa::Filtered<CandDstarWSelFlag>>>(cols, selectedCands, preslicSelectedCandDstarPerCol);
434+
}
435+
PROCESS_SWITCH(HfTaskDstarToD0Pi, processDataWoML, "Process Data without ML", true);
436+
437+
// process data function with using ML, Here we store BDT score as well
438+
void processDataWML(CollisionsWCent const& cols, soa::Filtered<CandDstarWSelFlagWMl> const& selectedCands)
439+
{
440+
runTaskDstar<true, soa::Filtered<CandDstarWSelFlagWMl>, Preslice<soa::Filtered<CandDstarWSelFlagWMl>>>(cols, selectedCands, preslicSelectedCandDstarPerColWMl);
441+
}
442+
PROCESS_SWITCH(HfTaskDstarToD0Pi, processDataWML, "Process Data with ML", false);
443+
444+
// process MC function without using ML
445+
void processMcWoMl(aod::McCollisions const&, CollisionsWCentMcLabel const& collisions, CandDstarWSelFlagMcRec const&,
446+
CandDstarMcGen const& rowsMcPartilces,
447+
aod::TracksWMc const&)
448+
{
449+
rowsSelectedCandDstarMcRec.bindExternalIndices(&collisions);
450+
runMcRecTaskDstar<false, Partition<CandDstarWSelFlagMcRec>>(rowsSelectedCandDstarMcRec, rowsMcPartilces);
451+
runMcGenTaskDstar(collisions, rowsMcPartilces);
452+
}
453+
PROCESS_SWITCH(HfTaskDstarToD0Pi, processMcWoMl, "Process MC Data without ML", false);
454+
455+
// process MC function with using ML
456+
void processMcWML(aod::McCollisions const&, CollisionsWCentMcLabel const& collisions, CandDstarWSelFlagWMlMcRec const&,
457+
CandDstarMcGen const& rowsMcPartilces,
458+
aod::TracksWMc const&)
459+
{
460+
rowsSelectedCandDstarMcRecWMl.bindExternalIndices(&collisions);
461+
runMcRecTaskDstar<true, Partition<CandDstarWSelFlagWMlMcRec>>(rowsSelectedCandDstarMcRecWMl, rowsMcPartilces);
462+
runMcGenTaskDstar(collisions, rowsMcPartilces);
414463
}
464+
PROCESS_SWITCH(HfTaskDstarToD0Pi, processMcWML, "Process MC Data with ML", false);
415465
};
416466

417467
WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)

0 commit comments

Comments
 (0)