Skip to content

Commit f4d05eb

Browse files
fgrosafgrosa
authored
Event filters: Add workflow to produce tables for BDT trainings (#174)
* Avoid to extract histo from registry by name * Event filters: Add workflow to produce tables for BDT trainings * Fix format
1 parent 081742b commit f4d05eb

File tree

1 file changed

+237
-15
lines changed

1 file changed

+237
-15
lines changed

‎EventFiltering/PWGHF/HFFilter.cxx‎

Lines changed: 237 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,8 @@
1313
/// \brief task for selection of events with HF signals
1414
///
1515
/// \author Fabrizio Grosa <fabrizio.grosa@cern.ch>, CERN
16+
/// \author Marcel Lesch <marcel.lesch@tum.de>, TUM
1617

17-
#include "Framework/runDataProcessing.h"
1818
#include "Framework/AnalysisDataModel.h"
1919
#include "Framework/AnalysisTask.h"
2020
#include "Framework/ASoAHelpers.h"
@@ -31,6 +31,23 @@
3131
#include "Math/Vector4D.h"
3232
#include "Math/GenVector/Boost.h"
3333

34+
using namespace o2;
35+
using namespace o2::framework;
36+
using namespace o2::framework::expressions;
37+
using namespace o2::aod::hf_cand;
38+
using namespace hf_cuts_single_track;
39+
40+
void customize(std::vector<o2::framework::ConfigParamSpec>& workflowOptions)
41+
{
42+
ConfigParamSpec optionTrigger{"doTrigger", VariantType::Bool, true, {"Perform trigger selection"}};
43+
ConfigParamSpec optionTraining{"doTrainSamples", VariantType::Bool, false, {"Produce training samples"}};
44+
workflowOptions.push_back(optionTrigger);
45+
workflowOptions.push_back(optionTraining);
46+
}
47+
48+
// always should be after customize() function
49+
#include "Framework/runDataProcessing.h"
50+
3451
namespace
3552
{
3653

@@ -48,6 +65,20 @@ enum BeautyCandType {
4865
kBeauty4Prong // combination of charm 3-prong and pion
4966
};
5067

68+
enum candOrig {
69+
kPrompt = 0,
70+
kNonPrompt,
71+
kBkg
72+
};
73+
74+
enum charmParticles {
75+
kD0 = 0,
76+
kDplus,
77+
kDs,
78+
kLc,
79+
kXic
80+
};
81+
5182
static const float massPi = RecoDecay::getMassPDG(211);
5283
static const float massK = RecoDecay::getMassPDG(321);
5384
static const float massProton = RecoDecay::getMassPDG(2212);
@@ -74,13 +105,95 @@ DECLARE_SOA_INDEX_COLUMN(Collision, collision);
74105
} // namespace extra3Prong
75106
DECLARE_SOA_TABLE(Colls2Prong, "AOD", "COLLSID2P", o2::aod::extra2Prong::CollisionId);
76107
DECLARE_SOA_TABLE(Colls3Prong, "AOD", "COLLSID3P", o2::aod::extra3Prong::CollisionId);
77-
} // namespace o2::aod
78108

79-
using namespace o2;
80-
using namespace o2::framework;
81-
using namespace o2::framework::expressions;
82-
using namespace o2::aod::hf_cand;
83-
using namespace hf_cuts_single_track;
109+
namespace hftraining2p
110+
{
111+
DECLARE_SOA_COLUMN(DCAPrimXY1, dcaPrimXY1, float); //!
112+
DECLARE_SOA_COLUMN(DCAPrimZ1, dcaPrimZ1, float); //!
113+
DECLARE_SOA_COLUMN(NsigmaPiTPC1, nsigmaPiTPC1, float); //!
114+
DECLARE_SOA_COLUMN(NsigmaKaTPC1, nsigmaKaTPC1, float); //!
115+
DECLARE_SOA_COLUMN(NsigmaPiTOF1, nsigmaPiTOF1, float); //!
116+
DECLARE_SOA_COLUMN(NsigmaKaTOF1, nsigmaKaTOF1, float); //!
117+
DECLARE_SOA_COLUMN(DCAPrimXY2, dcaPrimXY2, float); //!
118+
DECLARE_SOA_COLUMN(DCAPrimZ2, dcaPrimZ2, float); //!
119+
DECLARE_SOA_COLUMN(NsigmaPiTPC2, nsigmaPiTPC2, float); //!
120+
DECLARE_SOA_COLUMN(NsigmaKaTPC2, nsigmaKaTPC2, float); //!
121+
DECLARE_SOA_COLUMN(NsigmaPiTOF2, nsigmaPiTOF2, float); //!
122+
DECLARE_SOA_COLUMN(NsigmaKaTOF2, nsigmaKaTOF2, float); //!
123+
DECLARE_SOA_COLUMN(FlagOrigin, flagOrigin, int8_t); //!
124+
} // namespace hftraining2p
125+
DECLARE_SOA_TABLE(HFTrigTrain2P, "AOD", "HFTRIGTRAIN2P", //!
126+
hftraining2p::DCAPrimXY1,
127+
hftraining2p::DCAPrimZ1,
128+
hftraining2p::NsigmaPiTPC1,
129+
hftraining2p::NsigmaKaTPC1,
130+
hftraining2p::NsigmaPiTOF1,
131+
hftraining2p::NsigmaKaTOF1,
132+
hftraining2p::DCAPrimXY2,
133+
hftraining2p::DCAPrimZ2,
134+
hftraining2p::NsigmaPiTPC2,
135+
hftraining2p::NsigmaKaTPC2,
136+
hftraining2p::NsigmaPiTOF2,
137+
hftraining2p::NsigmaKaTOF2,
138+
hftraining2p::FlagOrigin);
139+
140+
namespace hftraining3p
141+
{
142+
DECLARE_SOA_COLUMN(DCAPrimXY1, dcaPrimXY1, float); //!
143+
DECLARE_SOA_COLUMN(DCAPrimZ1, dcaPrimZ1, float); //!
144+
DECLARE_SOA_COLUMN(NsigmaPiTPC1, nsigmaPiTPC1, float); //!
145+
DECLARE_SOA_COLUMN(NsigmaKaTPC1, nsigmaKaTPC1, float); //!
146+
DECLARE_SOA_COLUMN(NsigmaPrTPC1, nsigmaPrTPC1, float); //!
147+
DECLARE_SOA_COLUMN(NsigmaPiTOF1, nsigmaPiTOF1, float); //!
148+
DECLARE_SOA_COLUMN(NsigmaKaTOF1, nsigmaKaTOF1, float); //!
149+
DECLARE_SOA_COLUMN(NsigmaPrTOF1, nsigmaPrTOF1, float); //!
150+
DECLARE_SOA_COLUMN(DCAPrimXY2, dcaPrimXY2, float); //!
151+
DECLARE_SOA_COLUMN(DCAPrimZ2, dcaPrimZ2, float); //!
152+
DECLARE_SOA_COLUMN(NsigmaPiTPC2, nsigmaPiTPC2, float); //!
153+
DECLARE_SOA_COLUMN(NsigmaKaTPC2, nsigmaKaTPC2, float); //!
154+
DECLARE_SOA_COLUMN(NsigmaPrTPC2, nsigmaPrTPC2, float); //!
155+
DECLARE_SOA_COLUMN(NsigmaPiTOF2, nsigmaPiTOF2, float); //!
156+
DECLARE_SOA_COLUMN(NsigmaKaTOF2, nsigmaKaTOF2, float); //!
157+
DECLARE_SOA_COLUMN(NsigmaPrTOF2, nsigmaPrTOF2, float); //!
158+
DECLARE_SOA_COLUMN(DCAPrimXY3, dcaPrimXY3, float); //!
159+
DECLARE_SOA_COLUMN(DCAPrimZ3, dcaPrimZ3, float); //!
160+
DECLARE_SOA_COLUMN(NsigmaPiTPC3, nsigmaPiTPC3, float); //!
161+
DECLARE_SOA_COLUMN(NsigmaKaTPC3, nsigmaKaTPC3, float); //!
162+
DECLARE_SOA_COLUMN(NsigmaPrTPC3, nsigmaPrTPC3, float); //!
163+
DECLARE_SOA_COLUMN(NsigmaPiTOF3, nsigmaPiTOF3, float); //!
164+
DECLARE_SOA_COLUMN(NsigmaKaTOF3, nsigmaKaTOF3, float); //!
165+
DECLARE_SOA_COLUMN(NsigmaPrTOF3, nsigmaPrTOF3, float); //!
166+
DECLARE_SOA_COLUMN(FlagOrigin, flagOrigin, int8_t); //!
167+
DECLARE_SOA_COLUMN(Channel, channel, int8_t); //!
168+
} // namespace hftraining3p
169+
DECLARE_SOA_TABLE(HFTrigTrain3P, "AOD", "HFTRIGTRAIN3P", //!
170+
hftraining3p::DCAPrimXY1,
171+
hftraining3p::DCAPrimZ1,
172+
hftraining3p::NsigmaPiTPC1,
173+
hftraining3p::NsigmaKaTPC1,
174+
hftraining3p::NsigmaPrTPC1,
175+
hftraining3p::NsigmaPiTOF1,
176+
hftraining3p::NsigmaKaTOF1,
177+
hftraining3p::NsigmaPrTOF1,
178+
hftraining3p::DCAPrimXY2,
179+
hftraining3p::DCAPrimZ2,
180+
hftraining3p::NsigmaPiTPC2,
181+
hftraining3p::NsigmaKaTPC2,
182+
hftraining3p::NsigmaPrTPC2,
183+
hftraining3p::NsigmaPiTOF2,
184+
hftraining3p::NsigmaKaTOF2,
185+
hftraining3p::NsigmaPrTOF2,
186+
hftraining3p::DCAPrimXY3,
187+
hftraining3p::DCAPrimZ3,
188+
hftraining3p::NsigmaPiTPC3,
189+
hftraining3p::NsigmaKaTPC3,
190+
hftraining3p::NsigmaPrTPC3,
191+
hftraining3p::NsigmaPiTOF3,
192+
hftraining3p::NsigmaKaTOF3,
193+
hftraining3p::NsigmaPrTOF3,
194+
hftraining3p::FlagOrigin,
195+
hftraining3p::Channel);
196+
} // namespace o2::aod
84197

85198
struct AddCollisionId {
86199

@@ -100,7 +213,7 @@ struct AddCollisionId {
100213
}
101214
};
102215

103-
struct HfFilter {
216+
struct HfFilter { // Main struct for HF triggers
104217

105218
Produces<aod::HfFilters> tags;
106219

@@ -125,16 +238,17 @@ struct HfFilter {
125238
std::array<LabeledArray<double>, 2> cutsSingleTrackBeauty;
126239

127240
HistogramRegistry registry{"registry", {}, OutputObjHandlingPolicy::AnalysisObject, true, true};
241+
std::shared_ptr<TH1> hProcessedEvents;
128242

129243
void init(o2::framework::InitContext&)
130244
{
131245

132246
cutsSingleTrackBeauty = {cutsTrackBeauty3Prong, cutsTrackBeauty4Prong};
133247

134-
registry.add("fProcessedEvents", "HF - event filtered;;events", HistType::kTH1F, {{5, -0.5, 4.5}});
248+
hProcessedEvents = std::get<std::shared_ptr<TH1>>(registry.add("fProcessedEvents", "HF - event filtered;;events", HistType::kTH1F, {{5, -0.5, 4.5}}));
135249
std::array<std::string, 5> eventTitles = {"all", "rejected", "w/ high-#it{p}_{T} candidate", "w/ beauty candidate", "w/ femto candidate"};
136250
for (size_t iBin = 0; iBin < eventTitles.size(); iBin++) {
137-
registry.get<TH1>(HIST("fProcessedEvents"))->GetXaxis()->SetBinLabel(iBin + 1, eventTitles[iBin].data());
251+
hProcessedEvents->GetXaxis()->SetBinLabel(iBin + 1, eventTitles[iBin].data());
138252
}
139253
}
140254

@@ -344,20 +458,128 @@ struct HfFilter {
344458
tags(keepEvent[kHighPt], keepEvent[kBeauty], keepEvent[kFemto]);
345459

346460
if (!keepEvent[kHighPt] && !keepEvent[kBeauty] && !keepEvent[kFemto]) {
347-
registry.get<TH1>(HIST("fProcessedEvents"))->Fill(1);
461+
hProcessedEvents->Fill(1);
348462
} else {
349463
for (int iTrigger{0}; iTrigger < kNtriggersHF; iTrigger++) {
350464
if (keepEvent[iTrigger]) {
351-
registry.get<TH1>(HIST("fProcessedEvents"))->Fill(iTrigger + 2);
465+
hProcessedEvents->Fill(iTrigger + 2);
352466
}
353467
}
354468
}
355469
}
356470
};
357471

472+
struct ProduceTrainingSamples { // Struct for training samples
473+
474+
Produces<aod::HFTrigTrain2P> train2P;
475+
Produces<aod::HFTrigTrain3P> train3P;
476+
477+
using BigTracksMCPID = soa::Join<aod::BigTracksPID, aod::BigTracksMC>;
478+
479+
void process(aod::HfTrackIndexProng2 const& cand2Prongs,
480+
aod::HfTrackIndexProng3 const& cand3Prongs,
481+
aod::McParticles const& particlesMC,
482+
BigTracksMCPID const&)
483+
{
484+
for (const auto& cand2Prong : cand2Prongs) { // start loop over 2 prongs
485+
486+
auto trackPos = cand2Prong.index0_as<BigTracksMCPID>(); // positive daughter
487+
auto trackNeg = cand2Prong.index1_as<BigTracksMCPID>(); // negative daughter
488+
489+
int8_t sign = 0;
490+
int8_t flag = 0;
491+
int8_t origin = 0;
492+
493+
// D0(bar) → π± K∓
494+
auto indexRec = RecoDecay::getMatchedMCRec(particlesMC, std::array{trackPos, trackNeg}, pdg::Code::kD0, array{+kPiPlus, -kKPlus}, true, &sign);
495+
if (indexRec > -1) {
496+
auto particle = particlesMC.iteratorAt(indexRec);
497+
origin = (RecoDecay::getMother(particlesMC, particle, kBottom, true) > -1 ? OriginType::NonPrompt : OriginType::Prompt);
498+
if (origin == OriginType::NonPrompt) {
499+
flag = kNonPrompt;
500+
} else {
501+
flag = kPrompt;
502+
}
503+
} else {
504+
flag = kBkg;
505+
}
506+
507+
train2P(trackPos.dcaPrim0(), trackPos.dcaPrim1(), trackPos.tpcNSigmaPi(), trackPos.tpcNSigmaKa(), trackPos.tofNSigmaPi(), trackPos.tofNSigmaKa(),
508+
trackNeg.dcaPrim0(), trackNeg.dcaPrim1(), trackNeg.tpcNSigmaPi(), trackNeg.tpcNSigmaKa(), trackNeg.tofNSigmaPi(), trackNeg.tofNSigmaKa(),
509+
flag);
510+
} // end loop over 2-prong candidates
511+
512+
for (const auto& cand3Prong : cand3Prongs) { // start loop over 3 prongs
513+
514+
auto trackFirst = cand3Prong.index0_as<BigTracksMCPID>(); // first daughter
515+
auto trackSecond = cand3Prong.index1_as<BigTracksMCPID>(); // second daughter
516+
auto trackThird = cand3Prong.index2_as<BigTracksMCPID>(); // third daughter
517+
auto arrayDaughters = std::array{trackFirst, trackSecond, trackThird};
518+
519+
int8_t sign = 0;
520+
int8_t flag = 0;
521+
int8_t channel = 0;
522+
int8_t origin = 0;
523+
524+
// D± → π± K∓ π±
525+
auto indexRec = RecoDecay::getMatchedMCRec(particlesMC, arrayDaughters, pdg::Code::kDPlus, array{+kPiPlus, -kKPlus, +kPiPlus}, true, &sign);
526+
if (indexRec < 0) {
527+
// Ds± → K± K∓ π±
528+
indexRec = RecoDecay::getMatchedMCRec(particlesMC, arrayDaughters, 431, array{+kKPlus, -kKPlus, +kPiPlus}, true, &sign); //TODO: replace hard coded pdg code
529+
} else {
530+
channel = kDplus;
531+
}
532+
if (indexRec < 0) {
533+
// Λc± → p± K∓ π±
534+
indexRec = RecoDecay::getMatchedMCRec(particlesMC, arrayDaughters, pdg::Code::kLambdaCPlus, array{+kProton, -kKPlus, +kPiPlus}, true, &sign, 2);
535+
} else {
536+
channel = kDs;
537+
}
538+
if (indexRec < 0) {
539+
// Ξc± → p± K∓ π±
540+
indexRec = RecoDecay::getMatchedMCRec(particlesMC, arrayDaughters, pdg::Code::kXiCPlus, array{+kProton, -kKPlus, +kPiPlus}, true, &sign);
541+
if (indexRec > -1) {
542+
channel = kXic;
543+
}
544+
} else {
545+
channel = kLc;
546+
}
547+
548+
if (indexRec > -1) {
549+
auto particle = particlesMC.iteratorAt(indexRec);
550+
origin = (RecoDecay::getMother(particlesMC, particle, kBottom, true) > -1 ? OriginType::NonPrompt : OriginType::Prompt);
551+
if (origin == OriginType::NonPrompt) {
552+
flag = kNonPrompt;
553+
} else {
554+
flag = kPrompt;
555+
}
556+
} else {
557+
flag = kBkg;
558+
}
559+
560+
train3P(trackFirst.dcaPrim0(), trackFirst.dcaPrim1(), trackFirst.tpcNSigmaPi(), trackFirst.tpcNSigmaKa(), trackFirst.tpcNSigmaPr(), trackFirst.tofNSigmaPi(), trackFirst.tofNSigmaKa(), trackFirst.tofNSigmaPr(),
561+
trackSecond.dcaPrim0(), trackSecond.dcaPrim1(), trackSecond.tpcNSigmaPi(), trackSecond.tpcNSigmaKa(), trackSecond.tpcNSigmaPr(), trackSecond.tofNSigmaPi(), trackSecond.tofNSigmaKa(), trackSecond.tofNSigmaPr(),
562+
trackThird.dcaPrim0(), trackThird.dcaPrim1(), trackThird.tpcNSigmaPi(), trackThird.tpcNSigmaKa(), trackThird.tpcNSigmaPr(), trackThird.tofNSigmaPi(), trackThird.tofNSigmaKa(), trackThird.tofNSigmaPr(),
563+
flag, channel);
564+
} // end loop over 3-prong candidates
565+
}
566+
};
567+
358568
WorkflowSpec defineDataProcessing(ConfigContext const& cfg)
359569
{
360-
return WorkflowSpec{
361-
adaptAnalysisTask<AddCollisionId>(cfg),
362-
adaptAnalysisTask<HfFilter>(cfg)};
570+
571+
WorkflowSpec workflow{};
572+
573+
const bool doTrainSamples = cfg.options().get<bool>("doTrainSamples");
574+
if (doTrainSamples) {
575+
workflow.push_back(adaptAnalysisTask<ProduceTrainingSamples>(cfg));
576+
}
577+
578+
const bool doTrigger = cfg.options().get<bool>("doTrigger");
579+
if (doTrigger) {
580+
workflow.push_back(adaptAnalysisTask<AddCollisionId>(cfg));
581+
workflow.push_back(adaptAnalysisTask<HfFilter>(cfg));
582+
}
583+
584+
return workflow;
363585
}

0 commit comments

Comments
 (0)