Skip to content

Commit be41d49

Browse files
author
Valerio Di Bella
committed
Adding : ML and skimming in correlator + other minor change
1 parent cd56171 commit be41d49

4 files changed

Lines changed: 287 additions & 54 deletions

File tree

‎PWGHF/HFC/DataModel/ReducedDMesonPairsTables.h‎

Lines changed: 26 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,8 @@
2323
#include <Framework/ASoA.h>
2424
#include <Framework/AnalysisDataModel.h>
2525

26+
#include <cstdint>
27+
2628
namespace o2::aod
2729
{
2830
DECLARE_SOA_TABLE(HfCandDpFullEvs, "AOD", "HFCANDDPFULLEV",
@@ -38,6 +40,8 @@ DECLARE_SOA_TABLE(HfCandDpMcEvs, "AOD", "HFCANDDPMCEV",
3840

3941
namespace full
4042
{
43+
DECLARE_SOA_INDEX_COLUMN_FULL(HfCandDpFullEv, hfCandDpFullEv, int32_t, HfCandDpFullEvs, "");
44+
DECLARE_SOA_INDEX_COLUMN_FULL(HfCandDpMcEv, hfCandDpFullEvMC, int32_t, HfCandDpMcEvs, "");
4145
DECLARE_SOA_COLUMN(RSecondaryVertex, rSecondaryVertex, float); //! Radius of secondary vertex (cm)
4246
DECLARE_SOA_COLUMN(PtProng0, ptProng0, float); //! Transverse momentum of prong0 (GeV/c)
4347
DECLARE_SOA_COLUMN(PProng0, pProng0, float); //! Momentum of prong0 (GeV/c)
@@ -57,8 +61,6 @@ DECLARE_SOA_COLUMN(Eta, eta, float);
5761
DECLARE_SOA_COLUMN(Phi, phi, float); //! Azimuth angle of candidate
5862
DECLARE_SOA_COLUMN(E, e, float); //! Energy of candidate (GeV)
5963
DECLARE_SOA_COLUMN(Centrality, centrality, float); //! Collision centrality
60-
DECLARE_SOA_INDEX_COLUMN(HfCandDpMcEv, hfCandDpMcEv); //! The Mc collision index this MC particles belongs to
61-
DECLARE_SOA_INDEX_COLUMN(HfCandDpFullEv, hfCandDpFullEv); //! The collision index this candidate belongs to
6264
DECLARE_SOA_COLUMN(NSigTpcPi0, nSigTpcPi0, float); //! TPC Nsigma separation for prong0 with pion mass hypothesis
6365
DECLARE_SOA_COLUMN(NSigTpcKa0, nSigTpcKa0, float); //! TPC Nsigma separation for prong0 with kaon mass hypothesis
6466
DECLARE_SOA_COLUMN(NSigTofPi0, nSigTofPi0, float); //! TOF Nsigma separation for prong0 with pion mass hypothesis
@@ -71,12 +73,12 @@ DECLARE_SOA_COLUMN(NSigTpcPi2, nSigTpcPi2, float);
7173
DECLARE_SOA_COLUMN(NSigTpcKa2, nSigTpcKa2, float); //! TPC Nsigma separation for prong2 with kaon mass hypothesis
7274
DECLARE_SOA_COLUMN(NSigTofPi2, nSigTofPi2, float); //! TOF Nsigma separation for prong2 with pion mass hypothesis
7375
DECLARE_SOA_COLUMN(NSigTofKa2, nSigTofKa2, float); //! TOF Nsigma separation for prong2 with kaon mass hypothesis
74-
DECLARE_SOA_COLUMN(NSigTpcTofPi0, nSigTpcTofPi0, float); //! TPC and TOF combined Nsigma separation for prong0 with pion mass hypothesis
75-
DECLARE_SOA_COLUMN(NSigTpcTofKa0, nSigTpcTofKa0, float); //! TPC and TOF combined Nsigma separation for prong0 with kaon mass hypothesis
76-
DECLARE_SOA_COLUMN(NSigTpcTofPi1, nSigTpcTofPi1, float); //! TPC and TOF combined Nsigma separation for prong1 with pion mass hypothesis
77-
DECLARE_SOA_COLUMN(NSigTpcTofKa1, nSigTpcTofKa1, float); //! TPC and TOF combined Nsigma separation for prong1 with kaon mass hypothesis
78-
DECLARE_SOA_COLUMN(NSigTpcTofPi2, nSigTpcTofPi2, float); //! TPC and TOF combined Nsigma separation for prong2 with pion mass hypothesis
79-
DECLARE_SOA_COLUMN(NSigTpcTofKa2, nSigTpcTofKa2, float); //! TPC and TOF combined Nsigma separation for prong2 with kaon mass hypothesis
76+
DECLARE_SOA_COLUMN(TPCTofNSigmaPi0, tpcTofNSigmaPi0, float); //! TPC and TOF combined Nsigma separation for prong0 with pion mass hypothesis
77+
DECLARE_SOA_COLUMN(TPCTofNSigmaKa0, tpcTofNSigmaKa0, float); //! TPC and TOF combined Nsigma separation for prong0 with kaon mass hypothesis
78+
DECLARE_SOA_COLUMN(TPCTofNSigmaPi1, tpcTofNSigmaPi1, float); //! TPC and TOF combined Nsigma separation for prong1 with pion mass hypothesis
79+
DECLARE_SOA_COLUMN(TPCTofNSigmaKa1, tpcTofNSigmaKa1, float); //! TPC and TOF combined Nsigma separation for prong1 with kaon mass hypothesis
80+
DECLARE_SOA_COLUMN(TPCTofNSigmaPi2, tpcTofNSigmaPi2, float); //! TPC and TOF combined Nsigma separation for prong2 with pion mass hypothesis
81+
DECLARE_SOA_COLUMN(TPCTofNSigmaKa2, tpcTofNSigmaKa2, float); //! TPC and TOF combined Nsigma separation for prong2 with kaon mass hypothesis
8082
DECLARE_SOA_COLUMN(DecayLength, decayLength, float); //! Decay length of candidate (cm)
8183
DECLARE_SOA_COLUMN(DecayLengthXY, decayLengthXY, float); //! Transverse decay length of candidate (cm)
8284
DECLARE_SOA_COLUMN(DecayLengthNormalised, decayLengthNormalised, float); //! Normalised decay length of candidate
@@ -96,7 +98,7 @@ DECLARE_SOA_TABLE(HfCandDpMls, "AOD", "HFCANDDPML",
9698
full::MlScore0,
9799
full::MlScore1)
98100

99-
DECLARE_SOA_TABLE(HfCandDpTinys, "AOD", "HFCANDDPTINY",
101+
DECLARE_SOA_TABLE(HfCandDpTinys, "AOD", "HFCANDDPTINY", o2::soa::Index<>,
100102
full::CandidateSelFlag,
101103
full::M,
102104
full::Pt,
@@ -107,7 +109,7 @@ DECLARE_SOA_TABLE(HfCandDpTinys, "AOD", "HFCANDDPTINY",
107109
hf_cand_mc_flag::OriginMcRec,
108110
hf_cand_mc_flag::FlagMcDecayChanRec)
109111

110-
DECLARE_SOA_TABLE(HfCandDpLites, "AOD", "HFCANDDPLITE",
112+
DECLARE_SOA_TABLE(HfCandDpLites, "AOD", "HFCANDDPLITE", o2::soa::Index<>,
111113
hf_cand::Chi2PCA,
112114
full::DecayLength,
113115
full::DecayLengthXY,
@@ -126,20 +128,20 @@ DECLARE_SOA_TABLE(HfCandDpLites, "AOD", "HFCANDDPLITE",
126128
full::NSigTpcKa0,
127129
full::NSigTofPi0,
128130
full::NSigTofKa0,
129-
full::NSigTpcTofPi0,
130-
full::NSigTpcTofKa0,
131+
full::TPCTofNSigmaPi0,
132+
full::TPCTofNSigmaKa0,
131133
full::NSigTpcPi1,
132134
full::NSigTpcKa1,
133135
full::NSigTofPi1,
134136
full::NSigTofKa1,
135-
full::NSigTpcTofPi1,
136-
full::NSigTpcTofKa1,
137+
full::TPCTofNSigmaPi1,
138+
full::TPCTofNSigmaKa1,
137139
full::NSigTpcPi2,
138140
full::NSigTpcKa2,
139141
full::NSigTofPi2,
140142
full::NSigTofKa2,
141-
full::NSigTpcTofPi2,
142-
full::NSigTpcTofKa2,
143+
full::TPCTofNSigmaPi2,
144+
full::TPCTofNSigmaKa2,
143145
full::CandidateSelFlag,
144146
full::M,
145147
full::Pt,
@@ -155,7 +157,7 @@ DECLARE_SOA_TABLE(HfCandDpLites, "AOD", "HFCANDDPLITE",
155157
hf_cand_mc_flag::OriginMcRec,
156158
hf_cand_mc_flag::FlagMcDecayChanRec)
157159

158-
DECLARE_SOA_TABLE(HfCandDpFulls, "AOD", "HFCANDDPFULL",
160+
DECLARE_SOA_TABLE(HfCandDpFulls, "AOD", "HFCANDDPFULL", o2::soa::Index<>,
159161
hf_cand::XSecondaryVertex,
160162
hf_cand::YSecondaryVertex,
161163
hf_cand::ZSecondaryVertex,
@@ -201,20 +203,20 @@ DECLARE_SOA_TABLE(HfCandDpFulls, "AOD", "HFCANDDPFULL",
201203
full::NSigTpcKa0,
202204
full::NSigTofPi0,
203205
full::NSigTofKa0,
204-
full::NSigTpcTofPi0,
205-
full::NSigTpcTofKa0,
206+
full::TPCTofNSigmaPi0,
207+
full::TPCTofNSigmaKa0,
206208
full::NSigTpcPi1,
207209
full::NSigTpcKa1,
208210
full::NSigTofPi1,
209211
full::NSigTofKa1,
210-
full::NSigTpcTofPi1,
211-
full::NSigTpcTofKa1,
212+
full::TPCTofNSigmaPi1,
213+
full::TPCTofNSigmaKa1,
212214
full::NSigTpcPi2,
213215
full::NSigTpcKa2,
214216
full::NSigTofPi2,
215217
full::NSigTofKa2,
216-
full::NSigTpcTofPi2,
217-
full::NSigTpcTofKa2,
218+
full::TPCTofNSigmaPi2,
219+
full::TPCTofNSigmaKa2,
218220
full::CandidateSelFlag,
219221
full::M,
220222
full::Pt,
@@ -233,7 +235,7 @@ DECLARE_SOA_TABLE(HfCandDpFulls, "AOD", "HFCANDDPFULL",
233235
hf_cand_mc_flag::OriginMcRec,
234236
hf_cand_mc_flag::FlagMcDecayChanRec);
235237

236-
DECLARE_SOA_TABLE(HfCandDpMcPs, "AOD", "HFCANDDPMCP",
238+
DECLARE_SOA_TABLE(HfCandDpMcPs, "AOD", "HFCANDDPMCP", o2::soa::Index<>,
237239
full::Pt,
238240
full::Eta,
239241
full::Phi,

‎PWGHF/HFC/TableProducer/CMakeLists.txt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ o2physics_add_dpl_workflow(correlator-dplus-dminus
3636

3737
o2physics_add_dpl_workflow(correlator-dplus-dplus-reduced
3838
SOURCES correlatorDplusDplusReduced.cxx
39-
PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2Physics::EventFilteringUtils
39+
PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2Physics::EventFilteringUtils O2Physics::MLCore
4040
COMPONENT_NAME Analysis)
4141

4242
o2physics_add_dpl_workflow(correlator-dplus-hadrons

‎PWGHF/HFC/TableProducer/correlatorDplusDplusReduced.cxx‎

Lines changed: 123 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,8 @@
2020
#include "PWGHF/Core/CentralityEstimation.h"
2121
#include "PWGHF/Core/DecayChannels.h"
2222
#include "PWGHF/Core/HfHelper.h"
23+
#include "PWGHF/Core/HfMlResponseDplusToPiKPi.h"
24+
#include "PWGHF/Core/SelectorCuts.h"
2325
#include "PWGHF/DataModel/CandidateReconstructionTables.h"
2426
#include "PWGHF/DataModel/CandidateSelectionTables.h"
2527
#include "PWGHF/HFC/DataModel/ReducedDMesonPairsTables.h"
@@ -30,20 +32,24 @@
3032
#include "Common/DataModel/Centrality.h"
3133

3234
#include <CCDB/BasicCCDBManager.h>
35+
#include <CCDB/CcdbApi.h>
3336
#include <Framework/ASoA.h>
3437
#include <Framework/AnalysisDataModel.h>
3538
#include <Framework/AnalysisHelpers.h>
3639
#include <Framework/AnalysisTask.h>
40+
#include <Framework/Array2D.h>
3741
#include <Framework/Configurable.h>
3842
#include <Framework/Expressions.h>
3943
#include <Framework/HistogramRegistry.h>
4044
#include <Framework/InitContext.h>
4145
#include <Framework/runDataProcessing.h>
4246

4347
#include <cstdint>
48+
#include <cstdlib>
4449
#include <vector>
4550

4651
using namespace o2;
52+
using namespace o2::analysis;
4753
using namespace o2::framework;
4854
using namespace o2::framework::expressions;
4955
using namespace o2::hf_centrality;
@@ -69,6 +75,34 @@ struct HfCorrelatorDplusDplusReduced {
6975
Configurable<bool> cfgSkimmedProcessing{"cfgSkimmedProcessing", true, "Enables processing of skimmed datasets"};
7076
Configurable<bool> skipSingleD{"skipSingleD", true, "Skip collisions with one or less D candidates"};
7177

78+
Configurable<bool> applyMl{"applyMl", false, "Flag to apply ML selections"};
79+
Configurable<bool> applySkimming{"applySkimming", false, "Flag to apply Skimming selections"};
80+
Configurable<bool> loadModelsFromCCDB{"loadModelsFromCCDB", false, "Flag to enable or disable the loading of models from CCDB"};
81+
Configurable<std::vector<double>> binsPtMl{"binsPtMl", std::vector<double>{hf_cuts_ml::vecBinsPt}, "pT bin limits for ML application"};
82+
Configurable<std::vector<int>> cutDirMl{"cutDirMl", std::vector<int>{hf_cuts_ml::vecCutDir}, "Whether to reject score values greater or smaller than the threshold"};
83+
Configurable<LabeledArray<double>> cutsMl{"cutsMl", {hf_cuts_ml::Cuts[0], hf_cuts_ml::NBinsPt, hf_cuts_ml::NCutScores, hf_cuts_ml::labelsPt, hf_cuts_ml::labelsCutScore}, "ML selections per pT bin"};
84+
Configurable<int> nClassesMl{"nClassesMl", static_cast<int>(hf_cuts_ml::NCutScores), "Number of classes in ML model"};
85+
Configurable<std::string> ccdbUrl{"ccdbUrl", "http://alice-ccdb.cern.ch", "url of the ccdb repository"};
86+
Configurable<std::vector<std::string>> modelPathsCCDB{"modelPathsCCDB", std::vector<std::string>{"EventFiltering/PWGHF/BDTDPlus"}, "Paths of models on CCDB"};
87+
Configurable<std::vector<std::string>> onnxFileNames{"onnxFileNames", std::vector<std::string>{"ModelHandler_onnx_DPlusToKPiPi.onnx"}, "ONNX file names for each pT bin (if not from CCDB full path)"};
88+
Configurable<int64_t> timestampCCDB{"timestampCCDB", -1, "timestamp of the ONNX file for ML model used to query in CCDB"};
89+
Configurable<std::vector<std::string>> namesInputFeatures{"namesInputFeatures", std::vector<std::string>{"feature1", "feature2"}, "Names of ML model input features"};
90+
91+
Configurable<std::vector<double>> cutPtSkimming{"cutPtSkimming", {1, 5, 1000}, "pT bin limits for Skimming application"};
92+
Configurable<std::vector<double>> minM{"minM", {0.7, 0.7}, "Mass minimal for the cut for each pt bin"};
93+
Configurable<std::vector<double>> maxM{"maxM", {2.0, 2.1}, "Mass maximal for the cut for each pt bin"};
94+
Configurable<std::vector<double>> minCosTheta{"minCosTheta", {0.96, 0.98}, "CosTheta minimal for the cut for each pt bin"};
95+
Configurable<std::vector<double>> minDecayLength{"minDecayLength", {0.02, 0.03}, "DecayLength minimal for the cut for each pt bin"};
96+
Configurable<std::vector<double>> maxNsigmaTPC{"maxNsigmaTPC", {3, 3}, "NsigmaTPC maximal for the cut for each pt bin"};
97+
Configurable<std::vector<double>> maxNsigmaTOF{"maxNsigmaTOF", {3, 3}, "NsigmaTOF maximal for the cut for each pt bin"};
98+
99+
Configurable<std::vector<double>> binsPtSkimming{"binsPtSkimming", {0}, "pT bin limits for Skimming application"};
100+
101+
HfMlResponseDplusToPiKPi<float> hfMlResponse;
102+
103+
std::vector<float> outML;
104+
o2::ccdb::CcdbApi ccdbApi;
105+
72106
HfHelper hfHelper;
73107

74108
Service<o2::ccdb::BasicCCDBManager> ccdb;
@@ -102,6 +136,46 @@ struct HfCorrelatorDplusDplusReduced {
102136
if (cfgSkimmedProcessing) {
103137
zorroSummary.setObject(zorro.getZorroSummary());
104138
}
139+
140+
if (applyMl) {
141+
hfMlResponse.configure(binsPtMl, cutsMl, cutDirMl, nClassesMl);
142+
if (loadModelsFromCCDB) {
143+
ccdbApi.init(ccdbUrl);
144+
hfMlResponse.setModelPathsCCDB(onnxFileNames, ccdbApi, modelPathsCCDB, timestampCCDB);
145+
} else {
146+
hfMlResponse.setModelPathsLocal(onnxFileNames);
147+
}
148+
hfMlResponse.cacheInputFeaturesIndices(namesInputFeatures);
149+
hfMlResponse.init();
150+
}
151+
}
152+
153+
bool Skimming(auto candidate,
154+
const std::vector<double> PtcutSkimming,
155+
const std::vector<double> Mmin,
156+
const std::vector<double> Mmax,
157+
const std::vector<double> CosThetamin,
158+
const std::vector<double> DecayLengthmin,
159+
const std::vector<double> NsigmaTPCmax,
160+
const std::vector<double> NsigmaTOFmax)
161+
{
162+
if (candidate.pt() < PtcutSkimming[0] || candidate.pt() > PtcutSkimming[PtcutSkimming.size() - 1]) {
163+
return false;
164+
}
165+
for (long unsigned int i = 1; i < PtcutSkimming.size(); i++) {
166+
if (candidate.pt() <= PtcutSkimming[i]) {
167+
if (hfHelper.invMassDplusToPiKPi(candidate) < Mmin[i - 1] ||
168+
hfHelper.invMassDplusToPiKPi(candidate) > Mmax[i - 1] ||
169+
candidate.cpa() < CosThetamin[i - 1] ||
170+
candidate.decayLength() < DecayLengthmin[i - 1] ||
171+
candidate.nSigTofKa1() > NsigmaTOFmax[i - 1] ||
172+
candidate.nSigTpcKa1() > NsigmaTPCmax[i - 1]) {
173+
return false;
174+
}
175+
return true;
176+
}
177+
}
178+
return false;
105179
}
106180

107181
template <typename T>
@@ -127,14 +201,14 @@ struct HfCorrelatorDplusDplusReduced {
127201
channelMc = candidate.flagMcDecayChanRec();
128202
}
129203

130-
std::vector<float> outputMl = {-999., -999.};
204+
std::vector<float> outML = {-999., -999.};
131205
if constexpr (DoMl) {
132206
for (unsigned int iclass = 0; iclass < classMlIndexes->size(); iclass++) {
133-
outputMl[iclass] = candidate.mlProbDplusToPiKPi()[classMlIndexes->at(iclass)];
207+
outML[iclass] = candidate.mlProbDplusToPiKPi()[classMlIndexes->at(iclass)];
134208
}
135209
rowCandidateMl(
136-
outputMl[0],
137-
outputMl[1]);
210+
outML[0],
211+
outML[1]);
138212
}
139213

140214
float cent{-1.};
@@ -291,7 +365,9 @@ struct HfCorrelatorDplusDplusReduced {
291365
static int lastRunNumber = -1;
292366
// reserve memory
293367
rowCandidateFullEvents.reserve(collisions.size());
294-
if (fillCandidateLiteTable) {
368+
if (fillCandidateTinyTable) {
369+
rowCandidateTiny.reserve(candidates.size());
370+
} else if (fillCandidateLiteTable) {
295371
rowCandidateLite.reserve(candidates.size());
296372
} else {
297373
rowCandidateFull.reserve(candidates.size());
@@ -320,6 +396,26 @@ struct HfCorrelatorDplusDplusReduced {
320396
for (const auto& candidate : candidatesInThisCollision) {
321397
auto prongCandidate = candidate.prong1_as<aod::Tracks>();
322398
auto candidateSign = -prongCandidate.sign();
399+
400+
if (applySkimming &&
401+
!Skimming(candidate,
402+
cutPtSkimming,
403+
minM,
404+
maxM,
405+
minCosTheta,
406+
minDecayLength,
407+
maxNsigmaTPC,
408+
maxNsigmaTOF)) {
409+
continue;
410+
}
411+
412+
if (applyMl) {
413+
std::vector<float> inputFeatures = hfMlResponse.getInputFeatures(candidate);
414+
bool const isSelectedMl = hfMlResponse.isSelectedMl(inputFeatures, abs(candidate.pt()), outML);
415+
if (!isSelectedMl) {
416+
continue;
417+
}
418+
}
323419
fillCandidateTable<aod::Collisions>(candidate, rowCandidateFullEvents.lastIndex(), candidateSign);
324420
}
325421
}
@@ -332,7 +428,9 @@ struct HfCorrelatorDplusDplusReduced {
332428
{
333429
// reserve memory
334430
rowCandidateFullEvents.reserve(collisions.size());
335-
if (fillCandidateLiteTable) {
431+
if (fillCandidateTinyTable) {
432+
rowCandidateTiny.reserve(candidates.size());
433+
} else if (fillCandidateLiteTable) {
336434
rowCandidateLite.reserve(candidates.size());
337435
} else {
338436
rowCandidateFull.reserve(candidates.size());
@@ -348,6 +446,25 @@ struct HfCorrelatorDplusDplusReduced {
348446
for (const auto& candidate : candidatesInThisCollision) {
349447
auto prongCandidate = candidate.prong1_as<aod::Tracks>();
350448
auto candidateSign = -prongCandidate.sign();
449+
450+
if (applySkimming &&
451+
!Skimming(candidate,
452+
cutPtSkimming,
453+
minM,
454+
maxM,
455+
minCosTheta,
456+
minDecayLength,
457+
maxNsigmaTPC,
458+
maxNsigmaTOF)) {
459+
continue;
460+
}
461+
if (applyMl) {
462+
std::vector<float> inputFeatures = hfMlResponse.getInputFeatures(candidate);
463+
bool const isSelectedMl = hfMlResponse.isSelectedMl(inputFeatures, abs(candidate.pt()), outML);
464+
if (!isSelectedMl) {
465+
continue;
466+
}
467+
}
351468
fillCandidateTable<aod::Collisions, true>(candidate, rowCandidateFullEvents.lastIndex(), candidateSign);
352469
}
353470
}

0 commit comments

Comments
 (0)