Skip to content

Commit 0182a83

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

4 files changed

Lines changed: 288 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: 124 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,25 @@
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

47+
#include <string>
4348
#include <cstdint>
49+
#include <cstdlib>
4450
#include <vector>
4551

4652
using namespace o2;
53+
using namespace o2::analysis;
4754
using namespace o2::framework;
4855
using namespace o2::framework::expressions;
4956
using namespace o2::hf_centrality;
@@ -69,6 +76,34 @@ struct HfCorrelatorDplusDplusReduced {
6976
Configurable<bool> cfgSkimmedProcessing{"cfgSkimmedProcessing", true, "Enables processing of skimmed datasets"};
7077
Configurable<bool> skipSingleD{"skipSingleD", true, "Skip collisions with one or less D candidates"};
7178

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

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

107182
template <typename T>
@@ -127,14 +202,14 @@ struct HfCorrelatorDplusDplusReduced {
127202
channelMc = candidate.flagMcDecayChanRec();
128203
}
129204

130-
std::vector<float> outputMl = {-999., -999.};
205+
std::vector<float> outML = {-999., -999.};
131206
if constexpr (DoMl) {
132207
for (unsigned int iclass = 0; iclass < classMlIndexes->size(); iclass++) {
133-
outputMl[iclass] = candidate.mlProbDplusToPiKPi()[classMlIndexes->at(iclass)];
208+
outML[iclass] = candidate.mlProbDplusToPiKPi()[classMlIndexes->at(iclass)];
134209
}
135210
rowCandidateMl(
136-
outputMl[0],
137-
outputMl[1]);
211+
outML[0],
212+
outML[1]);
138213
}
139214

140215
float cent{-1.};
@@ -291,7 +366,9 @@ struct HfCorrelatorDplusDplusReduced {
291366
static int lastRunNumber = -1;
292367
// reserve memory
293368
rowCandidateFullEvents.reserve(collisions.size());
294-
if (fillCandidateLiteTable) {
369+
if (fillCandidateTinyTable) {
370+
rowCandidateTiny.reserve(candidates.size());
371+
} else if (fillCandidateLiteTable) {
295372
rowCandidateLite.reserve(candidates.size());
296373
} else {
297374
rowCandidateFull.reserve(candidates.size());
@@ -320,6 +397,26 @@ struct HfCorrelatorDplusDplusReduced {
320397
for (const auto& candidate : candidatesInThisCollision) {
321398
auto prongCandidate = candidate.prong1_as<aod::Tracks>();
322399
auto candidateSign = -prongCandidate.sign();
400+
401+
if (applySkimming &&
402+
!Skimming(candidate,
403+
cutPtSkimming,
404+
minM,
405+
maxM,
406+
minCosTheta,
407+
minDecayLength,
408+
maxNsigmaTPC,
409+
maxNsigmaTOF)) {
410+
continue;
411+
}
412+
413+
if (applyMl) {
414+
std::vector<float> inputFeatures = hfMlResponse.getInputFeatures(candidate);
415+
bool const isSelectedMl = hfMlResponse.isSelectedMl(inputFeatures, abs(candidate.pt()), outML);
416+
if (!isSelectedMl) {
417+
continue;
418+
}
419+
}
323420
fillCandidateTable<aod::Collisions>(candidate, rowCandidateFullEvents.lastIndex(), candidateSign);
324421
}
325422
}
@@ -332,7 +429,9 @@ struct HfCorrelatorDplusDplusReduced {
332429
{
333430
// reserve memory
334431
rowCandidateFullEvents.reserve(collisions.size());
335-
if (fillCandidateLiteTable) {
432+
if (fillCandidateTinyTable) {
433+
rowCandidateTiny.reserve(candidates.size());
434+
} else if (fillCandidateLiteTable) {
336435
rowCandidateLite.reserve(candidates.size());
337436
} else {
338437
rowCandidateFull.reserve(candidates.size());
@@ -348,6 +447,25 @@ struct HfCorrelatorDplusDplusReduced {
348447
for (const auto& candidate : candidatesInThisCollision) {
349448
auto prongCandidate = candidate.prong1_as<aod::Tracks>();
350449
auto candidateSign = -prongCandidate.sign();
450+
451+
if (applySkimming &&
452+
!Skimming(candidate,
453+
cutPtSkimming,
454+
minM,
455+
maxM,
456+
minCosTheta,
457+
minDecayLength,
458+
maxNsigmaTPC,
459+
maxNsigmaTOF)) {
460+
continue;
461+
}
462+
if (applyMl) {
463+
std::vector<float> inputFeatures = hfMlResponse.getInputFeatures(candidate);
464+
bool const isSelectedMl = hfMlResponse.isSelectedMl(inputFeatures, abs(candidate.pt()), outML);
465+
if (!isSelectedMl) {
466+
continue;
467+
}
468+
}
351469
fillCandidateTable<aod::Collisions, true>(candidate, rowCandidateFullEvents.lastIndex(), candidateSign);
352470
}
353471
}

0 commit comments

Comments
 (0)