Skip to content

Commit fbaab8f

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

4 files changed

Lines changed: 280 additions & 48 deletions

File tree

‎PWGHF/HFC/DataModel/ReducedDMesonPairsTables.h‎

Lines changed: 25 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222

2323
#include <Framework/ASoA.h>
2424
#include <Framework/AnalysisDataModel.h>
25+
#include <cstdint>
2526

2627
namespace o2::aod
2728
{
@@ -38,6 +39,8 @@ DECLARE_SOA_TABLE(HfCandDpMcEvs, "AOD", "HFCANDDPMCEV",
3839

3940
namespace full
4041
{
42+
DECLARE_SOA_INDEX_COLUMN_FULL(HfCandDpFullEv, hfCandDpFullEv, int32_t, HfCandDpFullEvs, "");
43+
DECLARE_SOA_INDEX_COLUMN_FULL(HfCandDpMcEv, hfCandDpFullEvMC, int32_t, HfCandDpMcEvs, "");
4144
DECLARE_SOA_COLUMN(RSecondaryVertex, rSecondaryVertex, float); //! Radius of secondary vertex (cm)
4245
DECLARE_SOA_COLUMN(PtProng0, ptProng0, float); //! Transverse momentum of prong0 (GeV/c)
4346
DECLARE_SOA_COLUMN(PProng0, pProng0, float); //! Momentum of prong0 (GeV/c)
@@ -57,8 +60,6 @@ DECLARE_SOA_COLUMN(Eta, eta, float);
5760
DECLARE_SOA_COLUMN(Phi, phi, float); //! Azimuth angle of candidate
5861
DECLARE_SOA_COLUMN(E, e, float); //! Energy of candidate (GeV)
5962
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
6263
DECLARE_SOA_COLUMN(NSigTpcPi0, nSigTpcPi0, float); //! TPC Nsigma separation for prong0 with pion mass hypothesis
6364
DECLARE_SOA_COLUMN(NSigTpcKa0, nSigTpcKa0, float); //! TPC Nsigma separation for prong0 with kaon mass hypothesis
6465
DECLARE_SOA_COLUMN(NSigTofPi0, nSigTofPi0, float); //! TOF Nsigma separation for prong0 with pion mass hypothesis
@@ -71,12 +72,12 @@ DECLARE_SOA_COLUMN(NSigTpcPi2, nSigTpcPi2, float);
7172
DECLARE_SOA_COLUMN(NSigTpcKa2, nSigTpcKa2, float); //! TPC Nsigma separation for prong2 with kaon mass hypothesis
7273
DECLARE_SOA_COLUMN(NSigTofPi2, nSigTofPi2, float); //! TOF Nsigma separation for prong2 with pion mass hypothesis
7374
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
75+
DECLARE_SOA_COLUMN(TPCTofNSigmaPi0, tpcTofNSigmaPi0, float); //! TPC and TOF combined Nsigma separation for prong0 with pion mass hypothesis
76+
DECLARE_SOA_COLUMN(TPCTofNSigmaKa0, tpcTofNSigmaKa0, float); //! TPC and TOF combined Nsigma separation for prong0 with kaon mass hypothesis
77+
DECLARE_SOA_COLUMN(TPCTofNSigmaPi1, tpcTofNSigmaPi1, float); //! TPC and TOF combined Nsigma separation for prong1 with pion mass hypothesis
78+
DECLARE_SOA_COLUMN(TPCTofNSigmaKa1, tpcTofNSigmaKa1, float); //! TPC and TOF combined Nsigma separation for prong1 with kaon mass hypothesis
79+
DECLARE_SOA_COLUMN(TPCTofNSigmaPi2, tpcTofNSigmaPi2, float); //! TPC and TOF combined Nsigma separation for prong2 with pion mass hypothesis
80+
DECLARE_SOA_COLUMN(TPCTofNSigmaKa2, tpcTofNSigmaKa2, float); //! TPC and TOF combined Nsigma separation for prong2 with kaon mass hypothesis
8081
DECLARE_SOA_COLUMN(DecayLength, decayLength, float); //! Decay length of candidate (cm)
8182
DECLARE_SOA_COLUMN(DecayLengthXY, decayLengthXY, float); //! Transverse decay length of candidate (cm)
8283
DECLARE_SOA_COLUMN(DecayLengthNormalised, decayLengthNormalised, float); //! Normalised decay length of candidate
@@ -96,7 +97,7 @@ DECLARE_SOA_TABLE(HfCandDpMls, "AOD", "HFCANDDPML",
9697
full::MlScore0,
9798
full::MlScore1)
9899

99-
DECLARE_SOA_TABLE(HfCandDpTinys, "AOD", "HFCANDDPTINY",
100+
DECLARE_SOA_TABLE(HfCandDpTinys, "AOD", "HFCANDDPTINY", o2::soa::Index<>,
100101
full::CandidateSelFlag,
101102
full::M,
102103
full::Pt,
@@ -107,7 +108,7 @@ DECLARE_SOA_TABLE(HfCandDpTinys, "AOD", "HFCANDDPTINY",
107108
hf_cand_mc_flag::OriginMcRec,
108109
hf_cand_mc_flag::FlagMcDecayChanRec)
109110

110-
DECLARE_SOA_TABLE(HfCandDpLites, "AOD", "HFCANDDPLITE",
111+
DECLARE_SOA_TABLE(HfCandDpLites, "AOD", "HFCANDDPLITE", o2::soa::Index<>,
111112
hf_cand::Chi2PCA,
112113
full::DecayLength,
113114
full::DecayLengthXY,
@@ -126,20 +127,20 @@ DECLARE_SOA_TABLE(HfCandDpLites, "AOD", "HFCANDDPLITE",
126127
full::NSigTpcKa0,
127128
full::NSigTofPi0,
128129
full::NSigTofKa0,
129-
full::NSigTpcTofPi0,
130-
full::NSigTpcTofKa0,
130+
full::TPCTofNSigmaPi0,
131+
full::TPCTofNSigmaKa0,
131132
full::NSigTpcPi1,
132133
full::NSigTpcKa1,
133134
full::NSigTofPi1,
134135
full::NSigTofKa1,
135-
full::NSigTpcTofPi1,
136-
full::NSigTpcTofKa1,
136+
full::TPCTofNSigmaPi1,
137+
full::TPCTofNSigmaKa1,
137138
full::NSigTpcPi2,
138139
full::NSigTpcKa2,
139140
full::NSigTofPi2,
140141
full::NSigTofKa2,
141-
full::NSigTpcTofPi2,
142-
full::NSigTpcTofKa2,
142+
full::TPCTofNSigmaPi2,
143+
full::TPCTofNSigmaKa2,
143144
full::CandidateSelFlag,
144145
full::M,
145146
full::Pt,
@@ -155,7 +156,7 @@ DECLARE_SOA_TABLE(HfCandDpLites, "AOD", "HFCANDDPLITE",
155156
hf_cand_mc_flag::OriginMcRec,
156157
hf_cand_mc_flag::FlagMcDecayChanRec)
157158

158-
DECLARE_SOA_TABLE(HfCandDpFulls, "AOD", "HFCANDDPFULL",
159+
DECLARE_SOA_TABLE(HfCandDpFulls, "AOD", "HFCANDDPFULL", o2::soa::Index<>,
159160
hf_cand::XSecondaryVertex,
160161
hf_cand::YSecondaryVertex,
161162
hf_cand::ZSecondaryVertex,
@@ -201,20 +202,20 @@ DECLARE_SOA_TABLE(HfCandDpFulls, "AOD", "HFCANDDPFULL",
201202
full::NSigTpcKa0,
202203
full::NSigTofPi0,
203204
full::NSigTofKa0,
204-
full::NSigTpcTofPi0,
205-
full::NSigTpcTofKa0,
205+
full::TPCTofNSigmaPi0,
206+
full::TPCTofNSigmaKa0,
206207
full::NSigTpcPi1,
207208
full::NSigTpcKa1,
208209
full::NSigTofPi1,
209210
full::NSigTofKa1,
210-
full::NSigTpcTofPi1,
211-
full::NSigTpcTofKa1,
211+
full::TPCTofNSigmaPi1,
212+
full::TPCTofNSigmaKa1,
212213
full::NSigTpcPi2,
213214
full::NSigTpcKa2,
214215
full::NSigTofPi2,
215216
full::NSigTofKa2,
216-
full::NSigTpcTofPi2,
217-
full::NSigTpcTofKa2,
217+
full::TPCTofNSigmaPi2,
218+
full::TPCTofNSigmaKa2,
218219
full::CandidateSelFlag,
219220
full::M,
220221
full::Pt,
@@ -233,7 +234,7 @@ DECLARE_SOA_TABLE(HfCandDpFulls, "AOD", "HFCANDDPFULL",
233234
hf_cand_mc_flag::OriginMcRec,
234235
hf_cand_mc_flag::FlagMcDecayChanRec);
235236

236-
DECLARE_SOA_TABLE(HfCandDpMcPs, "AOD", "HFCANDDPMCP",
237+
DECLARE_SOA_TABLE(HfCandDpMcPs, "AOD", "HFCANDDPMCP", o2::soa::Index<>,
237238
full::Pt,
238239
full::Eta,
239240
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: 119 additions & 2 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> outputMl;
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+
std::vector<double> PtcutSkimming,
155+
std::vector<double> Mmin,
156+
std::vector<double> Mmax,
157+
std::vector<double> CosThetamin,
158+
std::vector<double> DecayLengthmin,
159+
std::vector<double> NsigmaTPCmax,
160+
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>
@@ -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()), outputMl);
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()), outputMl);
464+
if (!isSelectedMl) {
465+
continue;
466+
}
467+
}
351468
fillCandidateTable<aod::Collisions, true>(candidate, rowCandidateFullEvents.lastIndex(), candidateSign);
352469
}
353470
}

0 commit comments

Comments
 (0)