Skip to content

Commit 17a99ea

Browse files
MarcellocostiMarcellocosti
committed
Separated process function with centrality
1 parent a95ffe4 commit 17a99ea

File tree

1 file changed

+128
-13
lines changed

1 file changed

+128
-13
lines changed

‎PWGHF/TableProducer/treeCreatorDplusToPiKPi.cxx‎

Lines changed: 128 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -257,6 +257,7 @@ struct HfTreeCreatorDplusToPiKPi {
257257
Configurable<float> downSampleBkgFactor{"downSampleBkgFactor", 1., "Fraction of background candidates to keep for ML trainings"};
258258
Configurable<float> ptMaxForDownSample{"ptMaxForDownSample", 10., "Maximum pt for the application of the downsampling factor"};
259259
Configurable<std::vector<int>> classMlIndexes{"classMlIndexes", {0, 2}, "Indexes of ML bkg and non-prompt scores."};
260+
Configurable<int> centEstimator{"centEstimator", 0, "Centrality estimation (None: 0, FT0C: 2, FT0M: 3)"};
260261

261262
HfHelper hfHelper;
262263

@@ -265,7 +266,7 @@ struct HfTreeCreatorDplusToPiKPi {
265266
using SelectedCandidatesMcWithMl = soa::Filtered<soa::Join<aod::HfCand3Prong, aod::HfCand3ProngMcRec, aod::HfSelDplusToPiKPi, aod::HfMlDplusToPiKPi>>;
266267
using TracksWPid = soa::Join<aod::Tracks, aod::TracksPidPi, aod::PidTpcTofFullPi, aod::TracksPidKa, aod::PidTpcTofFullKa>;
267268

268-
using CollisionsCent = soa::Join<aod::Collisions, aod::CentFT0Cs>;
269+
using CollisionsCent = soa::Join<aod::Collisions, aod::CentFT0Cs, aod::CentFT0Ms>;
269270

270271
Filter filterSelectCandidates = aod::hf_sel_candidate_dplus::isSelDplusToPiKPi >= selectionFlagDplus;
271272
Filter filterMcGenMatching = nabs(o2::aod::hf_cand_3prong::flagMcMatchGen) == static_cast<int8_t>(BIT(aod::hf_cand_3prong::DecayType::DplusToPiKPi));
@@ -291,7 +292,7 @@ struct HfTreeCreatorDplusToPiKPi {
291292
runNumber);
292293
}
293294

294-
template <bool doMc = false, bool doMl = false, typename T>
295+
template <typename Coll, bool doMc = false, bool doMl = false, typename T>
295296
void fillCandidateTable(const T& candidate)
296297
{
297298
int8_t flagMc = 0;
@@ -316,7 +317,12 @@ struct HfTreeCreatorDplusToPiKPi {
316317
auto prong0 = candidate.template prong0_as<TracksWPid>();
317318
auto prong1 = candidate.template prong1_as<TracksWPid>();
318319
auto prong2 = candidate.template prong2_as<TracksWPid>();
319-
auto coll = candidate.template collision_as<CollisionsCent>();
320+
321+
float cent{-1.};
322+
auto coll = candidate.template collision_as<Coll>();
323+
if (std::is_same_v<Coll, CollisionsCent> && centEstimator != CentralityEstimator::None) {
324+
cent = getCentralityColl(coll, centEstimator);
325+
}
320326

321327
if (fillCandidateLiteTable) {
322328
rowCandidateLite(
@@ -361,7 +367,7 @@ struct HfTreeCreatorDplusToPiKPi {
361367
candidate.eta(),
362368
candidate.phi(),
363369
hfHelper.yDplus(candidate),
364-
getCentralityColl(coll, CentralityEstimator::FT0C),
370+
cent,
365371
flagMc,
366372
originMc,
367373
channelMc);
@@ -443,14 +449,14 @@ struct HfTreeCreatorDplusToPiKPi {
443449
candidate.phi(),
444450
hfHelper.yDplus(candidate),
445451
hfHelper.eDplus(candidate),
446-
getCentralityColl(coll, CentralityEstimator::FT0C),
452+
cent,
447453
flagMc,
448454
originMc,
449455
channelMc);
450456
}
451457
}
452458

453-
void processData(CollisionsCent const& collisions,
459+
void processData(aod::Collisions const& collisions,
454460
soa::Filtered<soa::Join<aod::HfCand3Prong, aod::HfSelDplusToPiKPi>> const& candidates,
455461
TracksWPid const&)
456462
{
@@ -473,13 +479,13 @@ struct HfTreeCreatorDplusToPiKPi {
473479
continue;
474480
}
475481
}
476-
fillCandidateTable(candidate);
482+
fillCandidateTable<aod::Collisions>(candidate);
477483
}
478484
}
479485

480486
PROCESS_SWITCH(HfTreeCreatorDplusToPiKPi, processData, "Process data", true);
481487

482-
void processMc(CollisionsCent const& collisions,
488+
void processMc(aod::Collisions const& collisions,
483489
aod::McCollisions const&,
484490
SelectedCandidatesMc const& candidates,
485491
MatchedGenCandidatesMc const& particles,
@@ -500,7 +506,7 @@ struct HfTreeCreatorDplusToPiKPi {
500506
rowCandidateFull.reserve(reconstructedCandSig.size());
501507
}
502508
for (const auto& candidate : reconstructedCandSig) {
503-
fillCandidateTable<true>(candidate);
509+
fillCandidateTable<aod::Collisions, true>(candidate);
504510
}
505511
} else if (fillOnlySignalMl) {
506512
rowCandidateMl.reserve(reconstructedCandSigMl.size());
@@ -516,7 +522,7 @@ struct HfTreeCreatorDplusToPiKPi {
516522
continue;
517523
}
518524
}
519-
fillCandidateTable<true, true>(candidate);
525+
fillCandidateTable<aod::Collisions, true, true>(candidate);
520526
}
521527
} else if (fillOnlyBackground) {
522528
if (fillCandidateLiteTable) {
@@ -531,7 +537,7 @@ struct HfTreeCreatorDplusToPiKPi {
531537
continue;
532538
}
533539
}
534-
fillCandidateTable<true>(candidate);
540+
fillCandidateTable<aod::Collisions, true>(candidate);
535541
}
536542
} else {
537543
if (fillCandidateLiteTable) {
@@ -540,7 +546,7 @@ struct HfTreeCreatorDplusToPiKPi {
540546
rowCandidateFull.reserve(candidates.size());
541547
}
542548
for (const auto& candidate : candidates) {
543-
fillCandidateTable<true>(candidate);
549+
fillCandidateTable<aod::Collisions, true>(candidate);
544550
}
545551
}
546552

@@ -558,7 +564,116 @@ struct HfTreeCreatorDplusToPiKPi {
558564
}
559565
}
560566

561-
PROCESS_SWITCH(HfTreeCreatorDplusToPiKPi, processMc, "Process MC", false);
567+
void processDataWCent(CollisionsCent const& collisions,
568+
soa::Filtered<soa::Join<aod::HfCand3Prong, aod::HfSelDplusToPiKPi>> const& candidates,
569+
TracksWPid const&)
570+
{
571+
// Filling event properties
572+
rowCandidateFullEvents.reserve(collisions.size());
573+
for (const auto& collision : collisions) {
574+
fillEvent(collision, 0, 1);
575+
}
576+
577+
// Filling candidate properties
578+
if (fillCandidateLiteTable) {
579+
rowCandidateLite.reserve(candidates.size());
580+
} else {
581+
rowCandidateFull.reserve(candidates.size());
582+
}
583+
for (const auto& candidate : candidates) {
584+
if (downSampleBkgFactor < 1.) {
585+
float pseudoRndm = candidate.ptProng0() * 1000. - static_cast<int64_t>(candidate.ptProng0() * 1000);
586+
if (candidate.pt() < ptMaxForDownSample && pseudoRndm >= downSampleBkgFactor) {
587+
continue;
588+
}
589+
}
590+
fillCandidateTable<CollisionsCent>(candidate);
591+
}
592+
}
593+
594+
PROCESS_SWITCH(HfTreeCreatorDplusToPiKPi, processDataWCent, "Process data with cent", true);
595+
596+
void processMcWCent(CollisionsCent const& collisions,
597+
aod::McCollisions const&,
598+
SelectedCandidatesMc const& candidates,
599+
MatchedGenCandidatesMc const& particles,
600+
SelectedCandidatesMcWithMl const&,
601+
TracksWPid const&)
602+
{
603+
// Filling event properties
604+
rowCandidateFullEvents.reserve(collisions.size());
605+
for (const auto& collision : collisions) {
606+
fillEvent(collision, 0, 1);
607+
}
608+
609+
// Filling candidate properties
610+
if (fillOnlySignal) {
611+
if (fillCandidateLiteTable) {
612+
rowCandidateLite.reserve(reconstructedCandSig.size());
613+
} else {
614+
rowCandidateFull.reserve(reconstructedCandSig.size());
615+
}
616+
for (const auto& candidate : reconstructedCandSig) {
617+
fillCandidateTable<CollisionsCent, true>(candidate);
618+
}
619+
} else if (fillOnlySignalMl) {
620+
rowCandidateMl.reserve(reconstructedCandSigMl.size());
621+
if (fillCandidateLiteTable) {
622+
rowCandidateLite.reserve(reconstructedCandSigMl.size());
623+
} else {
624+
rowCandidateFull.reserve(reconstructedCandSigMl.size());
625+
}
626+
for (const auto& candidate : reconstructedCandSigMl) {
627+
if (downSampleBkgFactor < 1.) {
628+
float pseudoRndm = candidate.ptProng0() * 1000. - static_cast<int64_t>(candidate.ptProng0() * 1000);
629+
if (candidate.pt() < ptMaxForDownSample && pseudoRndm >= downSampleBkgFactor) {
630+
continue;
631+
}
632+
}
633+
fillCandidateTable<CollisionsCent, true, true>(candidate);
634+
}
635+
} else if (fillOnlyBackground) {
636+
if (fillCandidateLiteTable) {
637+
rowCandidateLite.reserve(reconstructedCandBkg.size());
638+
} else {
639+
rowCandidateFull.reserve(reconstructedCandBkg.size());
640+
}
641+
for (const auto& candidate : reconstructedCandBkg) {
642+
if (downSampleBkgFactor < 1.) {
643+
float pseudoRndm = candidate.ptProng0() * 1000. - static_cast<int64_t>(candidate.ptProng0() * 1000);
644+
if (candidate.pt() < ptMaxForDownSample && pseudoRndm >= downSampleBkgFactor) {
645+
continue;
646+
}
647+
}
648+
fillCandidateTable<CollisionsCent, true>(candidate);
649+
}
650+
} else {
651+
if (fillCandidateLiteTable) {
652+
rowCandidateLite.reserve(candidates.size());
653+
} else {
654+
rowCandidateFull.reserve(candidates.size());
655+
}
656+
for (const auto& candidate : candidates) {
657+
fillCandidateTable<CollisionsCent, true>(candidate);
658+
}
659+
}
660+
661+
// Filling particle properties
662+
rowCandidateFullParticles.reserve(particles.size());
663+
for (const auto& particle : particles) {
664+
rowCandidateFullParticles(
665+
particle.mcCollision().bcId(),
666+
particle.pt(),
667+
particle.eta(),
668+
particle.phi(),
669+
RecoDecay::y(particle.pVector(), o2::constants::physics::MassDPlus),
670+
particle.flagMcMatchGen(),
671+
particle.originMcGen());
672+
}
673+
}
674+
675+
676+
PROCESS_SWITCH(HfTreeCreatorDplusToPiKPi, processMcWCent, "Process MC with cent", false);
562677
};
563678

564679
WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)

0 commit comments

Comments
 (0)