Skip to content

Commit 1e12e46

Browse files
ataviragatavirag
authored
[PWGHF] Add ML variables in treeCreatorD0ToKPi (AliceO2Group#9761)
1 parent 7a51f81 commit 1e12e46

File tree

1 file changed

+142
-16
lines changed

1 file changed

+142
-16
lines changed

‎PWGHF/TableProducer/treeCreatorD0ToKPi.cxx

+142-16
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,12 @@ DECLARE_SOA_COLUMN(IsEventReject, isEventReject, int);
8181
DECLARE_SOA_COLUMN(RunNumber, runNumber, int);
8282
DECLARE_SOA_INDEX_COLUMN(McCollision, mcCollision);
8383
} // namespace full
84+
namespace ml
85+
{
86+
DECLARE_SOA_COLUMN(BdtOutputBkg, bdtOutputBkg, float);
87+
DECLARE_SOA_COLUMN(BdtOutputPrompt, bdtOutputPrompt, float);
88+
DECLARE_SOA_COLUMN(BdtOutputNonPrompt, bdtOutputNonPrompt, float);
89+
} // namespace ml
8490

8591
DECLARE_SOA_TABLE(HfCandD0Lites, "AOD", "HFCANDD0LITE",
8692
hf_cand::Chi2PCA,
@@ -201,6 +207,11 @@ DECLARE_SOA_TABLE(HfCandD0FullPs, "AOD", "HFCANDD0FULLP",
201207
full::OriginMcGen,
202208
full::McParticleId);
203209

210+
DECLARE_SOA_TABLE(HfCandD0Mls, "AOD", "HFCANDD0ML",
211+
ml::BdtOutputBkg,
212+
ml::BdtOutputNonPrompt,
213+
ml::BdtOutputPrompt);
214+
204215
} // namespace o2::aod
205216

206217
/// Writes the full information in an output TTree
@@ -209,6 +220,7 @@ struct HfTreeCreatorD0ToKPi {
209220
Produces<o2::aod::HfCandD0FullEvs> rowCandidateFullEvents;
210221
Produces<o2::aod::HfCandD0FullPs> rowCandidateFullParticles;
211222
Produces<o2::aod::HfCandD0Lites> rowCandidateLite;
223+
Produces<o2::aod::HfCandD0Mls> rowCandidateMl;
212224

213225
Configurable<bool> fillCandidateLiteTable{"fillCandidateLiteTable", false, "Switch to fill lite table with candidate properties"};
214226
// parameters for production of training samples
@@ -219,7 +231,9 @@ struct HfTreeCreatorD0ToKPi {
219231

220232
// using TracksWPid = soa::Join<aod::Tracks, aod::TracksPidPi, aod::PidTpcTofFullPi, aod::TracksPidKa, aod::PidTpcTofFullKa>;
221233
using SelectedCandidatesMc = soa::Filtered<soa::Join<aod::HfCand2ProngWPid, aod::HfCand2ProngMcRec, aod::HfSelD0>>;
234+
using SelectedCandidatesMcMl = soa::Filtered<soa::Join<aod::HfCand2ProngWPid, aod::HfCand2ProngMcRec, aod::HfSelD0, aod::HfMlD0>>;
222235
using SelectedCandidatesMcKf = soa::Filtered<soa::Join<aod::HfCand2ProngWPid, aod::HfCand2ProngKF, aod::HfCand2ProngMcRec, aod::HfSelD0>>;
236+
using SelectedCandidatesMcKfMl = soa::Filtered<soa::Join<aod::HfCand2ProngWPid, aod::HfCand2ProngKF, aod::HfCand2ProngMcRec, aod::HfSelD0, aod::HfMlD0>>;
223237
using MatchedGenCandidatesMc = soa::Filtered<soa::Join<aod::McParticles, aod::HfCand2ProngMcGen>>;
224238

225239
Filter filterSelectCandidates = aod::hf_sel_candidate_d0::isSelD0 >= 1 || aod::hf_sel_candidate_d0::isSelD0bar >= 1;
@@ -230,9 +244,17 @@ struct HfTreeCreatorD0ToKPi {
230244
Partition<SelectedCandidatesMcKf> reconstructedCandSigKF = nabs(aod::hf_cand_2prong::flagMcMatchRec) == static_cast<int8_t>(BIT(aod::hf_cand_2prong::DecayType::D0ToPiK));
231245
Partition<SelectedCandidatesMcKf> reconstructedCandBkgKF = nabs(aod::hf_cand_2prong::flagMcMatchRec) != static_cast<int8_t>(BIT(aod::hf_cand_2prong::DecayType::D0ToPiK));
232246

247+
Partition<SelectedCandidatesMcMl> reconstructedCandSigMl = nabs(aod::hf_cand_2prong::flagMcMatchRec) == static_cast<int8_t>(BIT(aod::hf_cand_2prong::DecayType::D0ToPiK));
248+
Partition<SelectedCandidatesMcMl> reconstructedCandBkgMl = nabs(aod::hf_cand_2prong::flagMcMatchRec) != static_cast<int8_t>(BIT(aod::hf_cand_2prong::DecayType::D0ToPiK));
249+
Partition<SelectedCandidatesMcKfMl> reconstructedCandSigKFMl = nabs(aod::hf_cand_2prong::flagMcMatchRec) == static_cast<int8_t>(BIT(aod::hf_cand_2prong::DecayType::D0ToPiK));
250+
Partition<SelectedCandidatesMcKfMl> reconstructedCandBkgKFMl = nabs(aod::hf_cand_2prong::flagMcMatchRec) != static_cast<int8_t>(BIT(aod::hf_cand_2prong::DecayType::D0ToPiK));
251+
233252
void init(InitContext const&)
234253
{
235-
std::array<bool, 8> doprocess{doprocessDataWithDCAFitterN, doprocessDataWithKFParticle, doprocessMcWithDCAFitterOnlySig, doprocessMcWithDCAFitterOnlyBkg, doprocessMcWithDCAFitterAll, doprocessMcWithKFParticleOnlySig, doprocessMcWithKFParticleOnlyBkg, doprocessMcWithKFParticleAll};
254+
std::array<bool, 16> doprocess{doprocessDataWithDCAFitterN, doprocessDataWithKFParticle, doprocessMcWithDCAFitterOnlySig, doprocessMcWithDCAFitterOnlyBkg,
255+
doprocessMcWithDCAFitterAll, doprocessMcWithKFParticleOnlySig, doprocessMcWithKFParticleOnlyBkg, doprocessMcWithKFParticleAll,
256+
doprocessDataWithDCAFitterNMl, doprocessDataWithKFParticleMl, doprocessMcWithDCAFitterOnlySigMl, doprocessMcWithDCAFitterOnlyBkgMl,
257+
doprocessMcWithDCAFitterAllMl, doprocessMcWithKFParticleOnlySigMl, doprocessMcWithKFParticleOnlyBkgMl, doprocessMcWithKFParticleAllMl};
236258
if (std::accumulate(doprocess.begin(), doprocess.end(), 0) != 1) {
237259
LOGP(fatal, "Only one process function can be enabled at a time.");
238260
}
@@ -251,7 +273,7 @@ struct HfTreeCreatorD0ToKPi {
251273
runNumber);
252274
}
253275

254-
template <typename T>
276+
template <bool applyMl, typename T>
255277
auto fillTable(const T& candidate, int candFlag, double invMass, double cosThetaStar, double topoChi2,
256278
double ct, double y, double e, int8_t flagMc, int8_t origin)
257279
{
@@ -356,9 +378,23 @@ struct HfTreeCreatorD0ToKPi {
356378
origin,
357379
candidate.globalIndex());
358380
}
381+
if constexpr (applyMl) {
382+
if (candidate.isSelD0()) {
383+
rowCandidateMl(
384+
candidate.mlProbD0()[0],
385+
candidate.mlProbD0()[1],
386+
candidate.mlProbD0()[2]);
387+
}
388+
if (candidate.isSelD0bar()) {
389+
rowCandidateMl(
390+
candidate.mlProbD0bar()[0],
391+
candidate.mlProbD0bar()[1],
392+
candidate.mlProbD0bar()[2]);
393+
}
394+
}
359395
}
360396

361-
template <int reconstructionType, typename CandType>
397+
template <int reconstructionType, bool applyMl, typename CandType>
362398
void processData(aod::Collisions const& collisions,
363399
CandType const& candidates,
364400
aod::Tracks const&, aod::BCs const&)
@@ -375,6 +411,9 @@ struct HfTreeCreatorD0ToKPi {
375411
} else {
376412
rowCandidateFull.reserve(candidates.size());
377413
}
414+
if constexpr (applyMl) {
415+
rowCandidateMl.reserve(candidates.size());
416+
}
378417
for (const auto& candidate : candidates) {
379418
if (downSampleBkgFactor < 1.) {
380419
float pseudoRndm = candidate.ptProng0() * 1000. - static_cast<int64_t>(candidate.ptProng0() * 1000);
@@ -396,10 +435,10 @@ struct HfTreeCreatorD0ToKPi {
396435
massD0bar = hfHelper.invMassD0barToKPi(candidate);
397436
}
398437
if (candidate.isSelD0()) {
399-
fillTable(candidate, 0, massD0, hfHelper.cosThetaStarD0(candidate), topolChi2PerNdf, ctD, yD, eD, 0, 0);
438+
fillTable<applyMl>(candidate, 0, massD0, hfHelper.cosThetaStarD0(candidate), topolChi2PerNdf, ctD, yD, eD, 0, 0);
400439
}
401440
if (candidate.isSelD0bar()) {
402-
fillTable(candidate, 1, massD0bar, hfHelper.cosThetaStarD0bar(candidate), topolChi2PerNdf, ctD, yD, eD, 0, 0);
441+
fillTable<applyMl>(candidate, 1, massD0bar, hfHelper.cosThetaStarD0bar(candidate), topolChi2PerNdf, ctD, yD, eD, 0, 0);
403442
}
404443
}
405444
}
@@ -409,20 +448,38 @@ struct HfTreeCreatorD0ToKPi {
409448
aod::Tracks const& tracks,
410449
aod::BCs const& bcs)
411450
{
412-
processData<aod::hf_cand::VertexerType::DCAFitter>(collisions, candidates, tracks, bcs);
451+
processData<aod::hf_cand::VertexerType::DCAFitter, false>(collisions, candidates, tracks, bcs);
413452
}
414453
PROCESS_SWITCH(HfTreeCreatorD0ToKPi, processDataWithDCAFitterN, "Process data with DCAFitterN", true);
415454

455+
void processDataWithDCAFitterNMl(aod::Collisions const& collisions,
456+
soa::Filtered<soa::Join<aod::HfCand2ProngWPid, aod::HfSelD0, aod::HfMlD0>> const& candidates,
457+
aod::Tracks const& tracks,
458+
aod::BCs const& bcs)
459+
{
460+
processData<aod::hf_cand::VertexerType::DCAFitter, true>(collisions, candidates, tracks, bcs);
461+
}
462+
PROCESS_SWITCH(HfTreeCreatorD0ToKPi, processDataWithDCAFitterNMl, "Process data with DCAFitterN and ML", false);
463+
416464
void processDataWithKFParticle(aod::Collisions const& collisions,
417465
soa::Filtered<soa::Join<aod::HfCand2ProngWPid, aod::HfCand2ProngKF, aod::HfSelD0>> const& candidates,
418466
aod::Tracks const& tracks,
419467
aod::BCs const& bcs)
420468
{
421-
processData<aod::hf_cand::VertexerType::KfParticle>(collisions, candidates, tracks, bcs);
469+
processData<aod::hf_cand::VertexerType::KfParticle, false>(collisions, candidates, tracks, bcs);
422470
}
423471
PROCESS_SWITCH(HfTreeCreatorD0ToKPi, processDataWithKFParticle, "Process data with KFParticle", false);
424472

425-
template <int reconstructionType, bool onlyBkg, bool onlySig, typename CandType>
473+
void processDataWithKFParticleMl(aod::Collisions const& collisions,
474+
soa::Filtered<soa::Join<aod::HfCand2ProngWPid, aod::HfCand2ProngKF, aod::HfSelD0, aod::HfMlD0>> const& candidates,
475+
aod::Tracks const& tracks,
476+
aod::BCs const& bcs)
477+
{
478+
processData<aod::hf_cand::VertexerType::KfParticle, true>(collisions, candidates, tracks, bcs);
479+
}
480+
PROCESS_SWITCH(HfTreeCreatorD0ToKPi, processDataWithKFParticleMl, "Process data with KFParticle and ML", false);
481+
482+
template <int reconstructionType, bool onlyBkg, bool onlySig, bool applyMl, typename CandType>
426483
void processMc(aod::Collisions const& collisions,
427484
aod::McCollisions const&,
428485
CandType const& candidates,
@@ -442,6 +499,9 @@ struct HfTreeCreatorD0ToKPi {
442499
} else {
443500
rowCandidateFull.reserve(candidates.size());
444501
}
502+
if constexpr (applyMl) {
503+
rowCandidateMl.reserve(candidates.size());
504+
}
445505
for (const auto& candidate : candidates) {
446506
if constexpr (onlyBkg) {
447507
if (TESTBIT(std::abs(candidate.flagMcMatchRec()), aod::hf_cand_2prong::DecayType::D0ToPiK)) {
@@ -473,10 +533,10 @@ struct HfTreeCreatorD0ToKPi {
473533
massD0bar = hfHelper.invMassD0barToKPi(candidate);
474534
}
475535
if (candidate.isSelD0()) {
476-
fillTable(candidate, 0, massD0, hfHelper.cosThetaStarD0(candidate), topolChi2PerNdf, ctD, yD, eD, candidate.flagMcMatchRec(), candidate.originMcRec());
536+
fillTable<applyMl>(candidate, 0, massD0, hfHelper.cosThetaStarD0(candidate), topolChi2PerNdf, ctD, yD, eD, candidate.flagMcMatchRec(), candidate.originMcRec());
477537
}
478538
if (candidate.isSelD0bar()) {
479-
fillTable(candidate, 1, massD0bar, hfHelper.cosThetaStarD0bar(candidate), topolChi2PerNdf, ctD, yD, eD, candidate.flagMcMatchRec(), candidate.originMcRec());
539+
fillTable<applyMl>(candidate, 1, massD0bar, hfHelper.cosThetaStarD0bar(candidate), topolChi2PerNdf, ctD, yD, eD, candidate.flagMcMatchRec(), candidate.originMcRec());
480540
}
481541
}
482542

@@ -504,64 +564,130 @@ struct HfTreeCreatorD0ToKPi {
504564
aod::Tracks const& tracks,
505565
aod::BCs const& bcs)
506566
{
507-
processMc<aod::hf_cand::VertexerType::DCAFitter, false, true>(collisions, mcCollisions, reconstructedCandSig, mcParticles, tracks, bcs);
567+
processMc<aod::hf_cand::VertexerType::DCAFitter, false, true, false>(collisions, mcCollisions, reconstructedCandSig, mcParticles, tracks, bcs);
508568
}
509569
PROCESS_SWITCH(HfTreeCreatorD0ToKPi, processMcWithDCAFitterOnlySig, "Process MC with DCAFitterN only for signals", false);
510570

571+
void processMcWithDCAFitterOnlySigMl(aod::Collisions const& collisions,
572+
aod::McCollisions const& mcCollisions,
573+
SelectedCandidatesMcMl const&,
574+
MatchedGenCandidatesMc const& mcParticles,
575+
aod::Tracks const& tracks,
576+
aod::BCs const& bcs)
577+
{
578+
processMc<aod::hf_cand::VertexerType::DCAFitter, false, true, true>(collisions, mcCollisions, reconstructedCandSigMl, mcParticles, tracks, bcs);
579+
}
580+
PROCESS_SWITCH(HfTreeCreatorD0ToKPi, processMcWithDCAFitterOnlySigMl, "Process MC with DCAFitterN only for signals and ML", false);
581+
511582
void processMcWithDCAFitterOnlyBkg(aod::Collisions const& collisions,
512583
aod::McCollisions const& mcCollisions,
513584
SelectedCandidatesMc const&,
514585
MatchedGenCandidatesMc const& mcParticles,
515586
aod::Tracks const& tracks,
516587
aod::BCs const& bcs)
517588
{
518-
processMc<aod::hf_cand::VertexerType::DCAFitter, true, false>(collisions, mcCollisions, reconstructedCandBkg, mcParticles, tracks, bcs);
589+
processMc<aod::hf_cand::VertexerType::DCAFitter, true, false, false>(collisions, mcCollisions, reconstructedCandBkg, mcParticles, tracks, bcs);
519590
}
520591
PROCESS_SWITCH(HfTreeCreatorD0ToKPi, processMcWithDCAFitterOnlyBkg, "Process MC with DCAFitterN only for background", false);
521592

593+
void processMcWithDCAFitterOnlyBkgMl(aod::Collisions const& collisions,
594+
aod::McCollisions const& mcCollisions,
595+
SelectedCandidatesMcMl const&,
596+
MatchedGenCandidatesMc const& mcParticles,
597+
aod::Tracks const& tracks,
598+
aod::BCs const& bcs)
599+
{
600+
processMc<aod::hf_cand::VertexerType::DCAFitter, true, false, true>(collisions, mcCollisions, reconstructedCandBkgMl, mcParticles, tracks, bcs);
601+
}
602+
PROCESS_SWITCH(HfTreeCreatorD0ToKPi, processMcWithDCAFitterOnlyBkgMl, "Process MC with DCAFitterN only for background with ML", false);
603+
522604
void processMcWithDCAFitterAll(aod::Collisions const& collisions,
523605
aod::McCollisions const& mcCollisions,
524606
SelectedCandidatesMc const& candidates,
525607
MatchedGenCandidatesMc const& mcParticles,
526608
aod::Tracks const& tracks,
527609
aod::BCs const& bcs)
528610
{
529-
processMc<aod::hf_cand::VertexerType::DCAFitter, false, false>(collisions, mcCollisions, candidates, mcParticles, tracks, bcs);
611+
processMc<aod::hf_cand::VertexerType::DCAFitter, false, false, false>(collisions, mcCollisions, candidates, mcParticles, tracks, bcs);
530612
}
531613
PROCESS_SWITCH(HfTreeCreatorD0ToKPi, processMcWithDCAFitterAll, "Process MC with DCAFitterN", false);
532614

615+
void processMcWithDCAFitterAllMl(aod::Collisions const& collisions,
616+
aod::McCollisions const& mcCollisions,
617+
SelectedCandidatesMcMl const& candidates,
618+
MatchedGenCandidatesMc const& mcParticles,
619+
aod::Tracks const& tracks,
620+
aod::BCs const& bcs)
621+
{
622+
processMc<aod::hf_cand::VertexerType::DCAFitter, false, false, true>(collisions, mcCollisions, candidates, mcParticles, tracks, bcs);
623+
}
624+
PROCESS_SWITCH(HfTreeCreatorD0ToKPi, processMcWithDCAFitterAllMl, "Process MC with DCAFitterN with ML", false);
625+
533626
void processMcWithKFParticleOnlySig(aod::Collisions const& collisions,
534627
aod::McCollisions const& mcCollisions,
535628
SelectedCandidatesMcKf const&,
536629
MatchedGenCandidatesMc const& mcParticles,
537630
aod::Tracks const& tracks,
538631
aod::BCs const& bcs)
539632
{
540-
processMc<aod::hf_cand::VertexerType::KfParticle, false, true>(collisions, mcCollisions, reconstructedCandSigKF, mcParticles, tracks, bcs);
633+
processMc<aod::hf_cand::VertexerType::KfParticle, false, true, false>(collisions, mcCollisions, reconstructedCandSigKF, mcParticles, tracks, bcs);
541634
}
542635
PROCESS_SWITCH(HfTreeCreatorD0ToKPi, processMcWithKFParticleOnlySig, "Process MC with KFParticle only for signals", false);
543636

637+
void processMcWithKFParticleOnlySigMl(aod::Collisions const& collisions,
638+
aod::McCollisions const& mcCollisions,
639+
SelectedCandidatesMcKfMl const&,
640+
MatchedGenCandidatesMc const& mcParticles,
641+
aod::Tracks const& tracks,
642+
aod::BCs const& bcs)
643+
{
644+
processMc<aod::hf_cand::VertexerType::KfParticle, false, true, true>(collisions, mcCollisions, reconstructedCandSigKFMl, mcParticles, tracks, bcs);
645+
}
646+
PROCESS_SWITCH(HfTreeCreatorD0ToKPi, processMcWithKFParticleOnlySigMl, "Process MC with KFParticle only for signals with ML", false);
647+
544648
void processMcWithKFParticleOnlyBkg(aod::Collisions const& collisions,
545649
aod::McCollisions const& mcCollisions,
546650
SelectedCandidatesMcKf const&,
547651
MatchedGenCandidatesMc const& mcParticles,
548652
aod::Tracks const& tracks,
549653
aod::BCs const& bcs)
550654
{
551-
processMc<aod::hf_cand::VertexerType::KfParticle, true, false>(collisions, mcCollisions, reconstructedCandBkgKF, mcParticles, tracks, bcs);
655+
processMc<aod::hf_cand::VertexerType::KfParticle, true, false, false>(collisions, mcCollisions, reconstructedCandBkgKF, mcParticles, tracks, bcs);
552656
}
553657
PROCESS_SWITCH(HfTreeCreatorD0ToKPi, processMcWithKFParticleOnlyBkg, "Process MC with KFParticle only for background", false);
554658

659+
void processMcWithKFParticleOnlyBkgMl(aod::Collisions const& collisions,
660+
aod::McCollisions const& mcCollisions,
661+
SelectedCandidatesMcKfMl const&,
662+
MatchedGenCandidatesMc const& mcParticles,
663+
aod::Tracks const& tracks,
664+
aod::BCs const& bcs)
665+
{
666+
processMc<aod::hf_cand::VertexerType::KfParticle, true, false, true>(collisions, mcCollisions, reconstructedCandBkgKFMl, mcParticles, tracks, bcs);
667+
}
668+
PROCESS_SWITCH(HfTreeCreatorD0ToKPi, processMcWithKFParticleOnlyBkgMl, "Process MC with KFParticle only for background with ML", false);
669+
555670
void processMcWithKFParticleAll(aod::Collisions const& collisions,
556671
aod::McCollisions const& mcCollisions,
557672
SelectedCandidatesMcKf const& candidates,
558673
MatchedGenCandidatesMc const& mcParticles,
559674
aod::Tracks const& tracks,
560675
aod::BCs const& bcs)
561676
{
562-
processMc<aod::hf_cand::VertexerType::KfParticle, false, false>(collisions, mcCollisions, candidates, mcParticles, tracks, bcs);
677+
processMc<aod::hf_cand::VertexerType::KfParticle, true, false, false>(collisions, mcCollisions, candidates, mcParticles, tracks, bcs);
563678
}
564679
PROCESS_SWITCH(HfTreeCreatorD0ToKPi, processMcWithKFParticleAll, "Process MC with KFParticle", false);
680+
681+
void processMcWithKFParticleAllMl(aod::Collisions const& collisions,
682+
aod::McCollisions const& mcCollisions,
683+
SelectedCandidatesMcKfMl const& candidates,
684+
MatchedGenCandidatesMc const& mcParticles,
685+
aod::Tracks const& tracks,
686+
aod::BCs const& bcs)
687+
{
688+
processMc<aod::hf_cand::VertexerType::KfParticle, false, false, true>(collisions, mcCollisions, candidates, mcParticles, tracks, bcs);
689+
}
690+
PROCESS_SWITCH(HfTreeCreatorD0ToKPi, processMcWithKFParticleAllMl, "Process MC with KFParticle with ML", false);
565691
};
566692

567693
WorkflowSpec defineDataProcessing(ConfigContext const& cfgc)

0 commit comments

Comments
 (0)