Skip to content

Commit ef70ab1

Browse files
ataviragalibuild
atavirag
and
alibuild
authored
[PWGHF] Add BDT variables to D0 MlResponse (#10082)
Co-authored-by: ALICE Action Bot <alibuild@cern.ch>
1 parent 3410087 commit ef70ab1

File tree

2 files changed

+41
-25
lines changed

2 files changed

+41
-25
lines changed

‎PWGHF/Core/HfMlResponseD0ToKPi.h‎

Lines changed: 29 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,22 @@
9999
break; \
100100
}
101101

102+
// Variation of CHECK_AND_FILL_VEC_D0_SIGNED(OBJECT, FEATURE, GETTER1, GETTER2)
103+
// where GETTER1 and GETTER2 are methods of the OBJECT, the variable
104+
// is filled depending on whether it is a D0 or a D0bar
105+
// and INDEX is the index of the vector
106+
#define CHECK_AND_FILL_VEC_D0_ML(OBJECT, FEATURE, GETTER1, GETTER2, INDEX) \
107+
case static_cast<uint8_t>(InputFeaturesD0ToKPi::FEATURE): { \
108+
if constexpr (usingMl) { \
109+
if (pdgCode == o2::constants::physics::kD0) { \
110+
inputFeatures.emplace_back(OBJECT.GETTER1()[INDEX]); \
111+
} else { \
112+
inputFeatures.emplace_back(OBJECT.GETTER2()[INDEX]); \
113+
} \
114+
} \
115+
break; \
116+
}
117+
102118
namespace o2::analysis
103119
{
104120
enum class InputFeaturesD0ToKPi : uint8_t {
@@ -140,6 +156,9 @@ enum class InputFeaturesD0ToKPi : uint8_t {
140156
nSigTpcTofKaExpKa,
141157
maxNormalisedDeltaIP,
142158
impactParameterProduct,
159+
bdtOutputBkg,
160+
bdtOutputNonPrompt,
161+
bdtOutputPrompt,
143162
cosThetaStar,
144163
cpa,
145164
cpaXY,
@@ -160,7 +179,7 @@ class HfMlResponseD0ToKPi : public HfMlResponse<TypeOutputScore>
160179
/// Method to get the input features vector needed for ML inference
161180
/// \param candidate is the D0 candidate
162181
/// \return inputFeatures vector
163-
template <typename T1>
182+
template <bool usingMl = false, typename T1>
164183
std::vector<float> getInputFeatures(T1 const& candidate, int const& pdgCode)
165184
{
166185
std::vector<float> inputFeatures;
@@ -183,10 +202,6 @@ class HfMlResponseD0ToKPi : public HfMlResponse<TypeOutputScore>
183202
CHECK_AND_FILL_VEC_D0_FULL(candidate, nSigTpcKa0, /*getter*/ nSigTpcKa0);
184203
CHECK_AND_FILL_VEC_D0_FULL(candidate, nSigTpcPi1, /*getter*/ nSigTpcPi1);
185204
CHECK_AND_FILL_VEC_D0_FULL(candidate, nSigTpcKa1, /*getter*/ nSigTpcKa1);
186-
// CHECK_AND_FILL_VEC_D0_OBJECT_HFHELPER_SIGNED(prong0, prong1, nSigTpcPiExpPi, tpcNSigmaPi);
187-
// CHECK_AND_FILL_VEC_D0_OBJECT_HFHELPER_SIGNED(prong0, prong1, nSigTpcKaExpPi, tpcNSigmaKa);
188-
// CHECK_AND_FILL_VEC_D0_OBJECT_HFHELPER_SIGNED(prong1, prong0, nSigTpcPiExpKa, tpcNSigmaPi);
189-
// CHECK_AND_FILL_VEC_D0_OBJECT_HFHELPER_SIGNED(prong1, prong0, nSigTpcKaExpKa, tpcNSigmaKa);
190205
CHECK_AND_FILL_VEC_D0_SIGNED(candidate, nSigTpcPiExpPi, nSigTpcPi0, nSigTpcPi1);
191206
CHECK_AND_FILL_VEC_D0_SIGNED(candidate, nSigTpcKaExpPi, nSigTpcKa0, nSigTpcKa1);
192207
CHECK_AND_FILL_VEC_D0_SIGNED(candidate, nSigTpcPiExpKa, nSigTpcPi1, nSigTpcPi0);
@@ -196,10 +211,6 @@ class HfMlResponseD0ToKPi : public HfMlResponse<TypeOutputScore>
196211
CHECK_AND_FILL_VEC_D0_FULL(candidate, nSigTofKa0, /*getter*/ nSigTofKa0);
197212
CHECK_AND_FILL_VEC_D0_FULL(candidate, nSigTofPi1, /*getter*/ nSigTofPi1);
198213
CHECK_AND_FILL_VEC_D0_FULL(candidate, nSigTofKa1, /*getter*/ nSigTofKa1);
199-
// CHECK_AND_FILL_VEC_D0_OBJECT_HFHELPER_SIGNED(prong0, prong1, nSigTofPiExpPi, tofNSigmaPi);
200-
// CHECK_AND_FILL_VEC_D0_OBJECT_HFHELPER_SIGNED(prong0, prong1, nSigTofKaExpPi, tofNSigmaKa);
201-
// CHECK_AND_FILL_VEC_D0_OBJECT_HFHELPER_SIGNED(prong1, prong0, nSigTofPiExpKa, tofNSigmaPi);
202-
// CHECK_AND_FILL_VEC_D0_OBJECT_HFHELPER_SIGNED(prong1, prong0, nSigTofKaExpKa, tofNSigmaKa);
203214
CHECK_AND_FILL_VEC_D0_SIGNED(candidate, nSigTofPiExpPi, nSigTofPi0, nSigTofPi1);
204215
CHECK_AND_FILL_VEC_D0_SIGNED(candidate, nSigTofKaExpPi, nSigTofKa0, nSigTofKa1);
205216
CHECK_AND_FILL_VEC_D0_SIGNED(candidate, nSigTofPiExpKa, nSigTofPi1, nSigTofPi0);
@@ -209,15 +220,15 @@ class HfMlResponseD0ToKPi : public HfMlResponse<TypeOutputScore>
209220
CHECK_AND_FILL_VEC_D0_FULL(candidate, nSigTpcTofKa0, tpcTofNSigmaKa0);
210221
CHECK_AND_FILL_VEC_D0_FULL(candidate, nSigTpcTofPi1, tpcTofNSigmaPi1);
211222
CHECK_AND_FILL_VEC_D0_FULL(candidate, nSigTpcTofKa1, tpcTofNSigmaKa1);
212-
// CHECK_AND_FILL_VEC_D0_OBJECT_HFHELPER_SIGNED(prong0, prong1, nSigTpcTofPiExpPi, tpcTofNSigmaPi);
213-
// CHECK_AND_FILL_VEC_D0_OBJECT_HFHELPER_SIGNED(prong0, prong1, nSigTpcTofKaExpPi, tpcTofNSigmaKa);
214-
// CHECK_AND_FILL_VEC_D0_OBJECT_HFHELPER_SIGNED(prong1, prong0, nSigTpcTofPiExpKa, tpcTofNSigmaPi);
215-
// CHECK_AND_FILL_VEC_D0_OBJECT_HFHELPER_SIGNED(prong1, prong0, nSigTpcTofKaExpKa, tpcTofNSigmaKa);
216223
CHECK_AND_FILL_VEC_D0_SIGNED(candidate, nSigTpcTofPiExpPi, tpcTofNSigmaPi0, tpcTofNSigmaPi1);
217224
CHECK_AND_FILL_VEC_D0_SIGNED(candidate, nSigTpcTofKaExpPi, tpcTofNSigmaKa0, tpcTofNSigmaKa1);
218225
CHECK_AND_FILL_VEC_D0_SIGNED(candidate, nSigTpcTofPiExpKa, tpcTofNSigmaPi1, tpcTofNSigmaPi0);
219226
CHECK_AND_FILL_VEC_D0_SIGNED(candidate, nSigTpcTofKaExpKa, tpcTofNSigmaKa1, tpcTofNSigmaKa0);
220227

228+
CHECK_AND_FILL_VEC_D0_ML(candidate, bdtOutputBkg, mlProbD0, mlProbD0bar, 0);
229+
CHECK_AND_FILL_VEC_D0_ML(candidate, bdtOutputNonPrompt, mlProbD0, mlProbD0bar, 1);
230+
CHECK_AND_FILL_VEC_D0_ML(candidate, bdtOutputPrompt, mlProbD0, mlProbD0bar, 2);
231+
221232
CHECK_AND_FILL_VEC_D0(maxNormalisedDeltaIP);
222233
CHECK_AND_FILL_VEC_D0_FULL(candidate, impactParameterProduct, impactParameterProduct);
223234
CHECK_AND_FILL_VEC_D0_HFHELPER_SIGNED(candidate, cosThetaStar, cosThetaStarD0, cosThetaStarD0bar);
@@ -273,6 +284,10 @@ class HfMlResponseD0ToKPi : public HfMlResponse<TypeOutputScore>
273284
FILL_MAP_D0(nSigTpcTofKaExpPi),
274285
FILL_MAP_D0(nSigTpcTofPiExpKa),
275286
FILL_MAP_D0(nSigTpcTofKaExpKa),
287+
// ML variables
288+
FILL_MAP_D0(bdtOutputBkg),
289+
FILL_MAP_D0(bdtOutputNonPrompt),
290+
FILL_MAP_D0(bdtOutputPrompt),
276291

277292
FILL_MAP_D0(maxNormalisedDeltaIP),
278293
FILL_MAP_D0(impactParameterProduct),
@@ -291,5 +306,6 @@ class HfMlResponseD0ToKPi : public HfMlResponse<TypeOutputScore>
291306
#undef CHECK_AND_FILL_VEC_D0_HFHELPER
292307
#undef CHECK_AND_FILL_VEC_D0_HFHELPER_SIGNED
293308
#undef CHECK_AND_FILL_VEC_D0_OBJECT_HFHELPER_SIGNED
309+
#undef CHECK_AND_FILL_VEC_D0_ML
294310

295311
#endif // PWGHF_CORE_HFMLRESPONSED0TOKPI_H_

‎PWGHF/HFC/TableProducer/correlatorDMesonPairs.cxx‎

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -105,8 +105,8 @@ struct HfCorrelatorDMesonPairs {
105105

106106
// using TracksWPid = soa::Join<aod::Tracks, aod::TracksPidPi, aod::PidTpcTofFullPi, aod::TracksPidKa, aod::PidTpcTofFullKa>;
107107

108-
Partition<soa::Join<aod::HfCand2ProngWPid, aod::HfSelD0>> selectedD0Candidates = aod::hf_sel_candidate_d0::isSelD0 >= selectionFlagD0 || aod::hf_sel_candidate_d0::isSelD0bar >= selectionFlagD0bar;
109-
Partition<soa::Join<aod::HfCand2ProngWPid, aod::HfSelD0, aod::HfCand2ProngMcRec>> selectedD0CandidatesMc = aod::hf_sel_candidate_d0::isRecoHfFlag >= selectionFlagHf;
108+
Partition<soa::Join<aod::HfCand2ProngWPid, aod::HfSelD0, aod::HfMlD0>> selectedD0Candidates = aod::hf_sel_candidate_d0::isSelD0 >= selectionFlagD0 || aod::hf_sel_candidate_d0::isSelD0bar >= selectionFlagD0bar;
109+
Partition<soa::Join<aod::HfCand2ProngWPid, aod::HfSelD0, aod::HfMlD0, aod::HfCand2ProngMcRec>> selectedD0CandidatesMc = aod::hf_sel_candidate_d0::isRecoHfFlag >= selectionFlagHf;
110110

111111
HistogramConfigSpec hTH1Pt{HistType::kTH1F, {{180, 0., 36.}}};
112112
HistogramConfigSpec hTH1Y{HistType::kTH1F, {{100, -5., 5.}}};
@@ -524,7 +524,7 @@ struct HfCorrelatorDMesonPairs {
524524

525525
/// D0(bar)-D0(bar) correlation pair builder - for real data and data-like analysis (i.e. reco-level w/o matching request via MC truth)
526526
void processData(aod::Collision const& collision,
527-
soa::Join<aod::HfCand2ProngWPid, aod::HfSelD0> const& candidates, aod::Tracks const&)
527+
soa::Join<aod::HfCand2ProngWPid, aod::HfSelD0, aod::HfMlD0> const& candidates, aod::Tracks const&)
528528
{
529529
for (const auto& candidate : candidates) {
530530
analysePid(candidate);
@@ -564,11 +564,11 @@ struct HfCorrelatorDMesonPairs {
564564

565565
if (applyMl) {
566566
if (isDCand1) {
567-
std::vector<float> inputFeaturesD0 = hfMlResponse.getInputFeatures(candidate1, o2::constants::physics::kD0);
567+
std::vector<float> inputFeaturesD0 = hfMlResponse.getInputFeatures<true>(candidate1, o2::constants::physics::kD0);
568568
isSelectedMlD0Cand1 = hfMlResponse.isSelectedMl(inputFeaturesD0, candidate1.pt(), outputMlD0Cand1);
569569
}
570570
if (isDbarCand1) {
571-
std::vector<float> inputFeaturesD0bar = hfMlResponse.getInputFeatures(candidate1, o2::constants::physics::kD0Bar);
571+
std::vector<float> inputFeaturesD0bar = hfMlResponse.getInputFeatures<true>(candidate1, o2::constants::physics::kD0Bar);
572572
isSelectedMlD0barCand1 = hfMlResponse.isSelectedMl(inputFeaturesD0bar, candidate1.pt(), outputMlD0barCand1);
573573
}
574574

@@ -630,11 +630,11 @@ struct HfCorrelatorDMesonPairs {
630630

631631
if (applyMl) {
632632
if (isDCand2) {
633-
std::vector<float> inputFeaturesD0 = hfMlResponse.getInputFeatures(candidate2, o2::constants::physics::kD0);
633+
std::vector<float> inputFeaturesD0 = hfMlResponse.getInputFeatures<true>(candidate2, o2::constants::physics::kD0);
634634
isSelectedMlD0Cand2 = hfMlResponse.isSelectedMl(inputFeaturesD0, candidate2.pt(), outputMlD0Cand2);
635635
}
636636
if (isDbarCand2) {
637-
std::vector<float> inputFeaturesD0bar = hfMlResponse.getInputFeatures(candidate2, o2::constants::physics::kD0Bar);
637+
std::vector<float> inputFeaturesD0bar = hfMlResponse.getInputFeatures<true>(candidate2, o2::constants::physics::kD0Bar);
638638
isSelectedMlD0barCand2 = hfMlResponse.isSelectedMl(inputFeaturesD0bar, candidate2.pt(), outputMlD0barCand2);
639639
}
640640

@@ -665,7 +665,7 @@ struct HfCorrelatorDMesonPairs {
665665

666666
PROCESS_SWITCH(HfCorrelatorDMesonPairs, processData, "Process data mode", true);
667667

668-
void processMcRec(aod::Collision const& collision, soa::Join<aod::HfCand2ProngWPid, aod::HfSelD0, aod::HfCand2ProngMcRec> const& candidates, aod::Tracks const&)
668+
void processMcRec(aod::Collision const& collision, soa::Join<aod::HfCand2ProngWPid, aod::HfSelD0, aod::HfMlD0, aod::HfCand2ProngMcRec> const& candidates, aod::Tracks const&)
669669
{
670670
for (const auto& candidate : candidates) {
671671
analysePid(candidate);
@@ -719,11 +719,11 @@ struct HfCorrelatorDMesonPairs {
719719

720720
if (applyMl) {
721721
if (isDCand1) {
722-
std::vector<float> inputFeaturesD0 = hfMlResponse.getInputFeatures(candidate1, o2::constants::physics::kD0);
722+
std::vector<float> inputFeaturesD0 = hfMlResponse.getInputFeatures<true>(candidate1, o2::constants::physics::kD0);
723723
isSelectedMlD0Cand1 = hfMlResponse.isSelectedMl(inputFeaturesD0, candidate1.pt(), outputMlD0Cand1);
724724
}
725725
if (isDbarCand1) {
726-
std::vector<float> inputFeaturesD0bar = hfMlResponse.getInputFeatures(candidate1, o2::constants::physics::kD0Bar);
726+
std::vector<float> inputFeaturesD0bar = hfMlResponse.getInputFeatures<true>(candidate1, o2::constants::physics::kD0Bar);
727727
isSelectedMlD0barCand1 = hfMlResponse.isSelectedMl(inputFeaturesD0bar, candidate1.pt(), outputMlD0barCand1);
728728
}
729729
// Remove non-ML selected D0 candidates
@@ -826,11 +826,11 @@ struct HfCorrelatorDMesonPairs {
826826

827827
if (applyMl) {
828828
if (isDCand2) {
829-
std::vector<float> inputFeaturesD0 = hfMlResponse.getInputFeatures(candidate2, o2::constants::physics::kD0);
829+
std::vector<float> inputFeaturesD0 = hfMlResponse.getInputFeatures<true>(candidate2, o2::constants::physics::kD0);
830830
isSelectedMlD0Cand2 = hfMlResponse.isSelectedMl(inputFeaturesD0, candidate2.pt(), outputMlD0Cand2);
831831
}
832832
if (isDbarCand2) {
833-
std::vector<float> inputFeaturesD0bar = hfMlResponse.getInputFeatures(candidate2, o2::constants::physics::kD0Bar);
833+
std::vector<float> inputFeaturesD0bar = hfMlResponse.getInputFeatures<true>(candidate2, o2::constants::physics::kD0Bar);
834834
isSelectedMlD0barCand2 = hfMlResponse.isSelectedMl(inputFeaturesD0bar, candidate2.pt(), outputMlD0barCand2);
835835
}
836836

0 commit comments

Comments
 (0)