Skip to content

Commit 9629f7a

Browse files
Tao-FangTao-Fang
authored
Add files via upload
1 parent 6fd6bb1 commit 9629f7a

File tree

3 files changed

+66
-9
lines changed

3 files changed

+66
-9
lines changed

‎PWGHF/TableProducer/CMakeLists.txt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -184,7 +184,7 @@ o2physics_add_dpl_workflow(candidate-selector-xic0-to-xi-pi-kf
184184

185185
o2physics_add_dpl_workflow(candidate-selector-to-xi-pi
186186
SOURCES candidateSelectorToXiPi.cxx
187-
PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore
187+
PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2Physics::MLCore
188188
COMPONENT_NAME Analysis)
189189

190190
o2physics_add_dpl_workflow(candidate-selector-xic-to-p-k-pi

‎PWGHF/TableProducer/candidateSelectorToXiPi.cxx‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,13 @@
1212
/// \file candidateSelectorToXiPi.cxx
1313
/// \brief Xic0 and Omegac0 → Xi Pi selection task
1414
/// \author Federica Zanone <federica.zanone@cern.ch>, Heidelberg University
15+
/// \author Tao Fang <tao.fang@cern.ch>, Central China Normal University
16+
17+
#include <string>
18+
#include <vector>
19+
20+
#include "PWGHF/Core/HfMlResponseXic0ToXiPi.h"
21+
#include "PWGHF/Core/SelectorCuts.h"
1522

1623
#include "PWGHF/DataModel/CandidateReconstructionTables.h"
1724
#include "PWGHF/DataModel/CandidateSelectionTables.h"
@@ -54,6 +61,7 @@ enum PidInfoStored {
5461
/// Struct for applying Omegac0/Xic0 selection cuts
5562
struct HfCandidateSelectorToXiPi {
5663
Produces<aod::HfSelToXiPi> hfSelToXiPi;
64+
Produces<aod::HfMlToXiPiKf> hfMlToXiPi;
5765

5866
// LF analysis selections
5967
Configurable<double> radiusCascMin{"radiusCascMin", 0.6, "Min cascade radius"};
@@ -129,6 +137,25 @@ struct HfCandidateSelectorToXiPi {
129137
Configurable<int> nClustersItsInnBarrMin{"nClustersItsInnBarrMin", 1, "Minimum number of ITS clusters in inner barrel requirement for pi <- charm baryon"};
130138
Configurable<float> itsChi2PerClusterMax{"itsChi2PerClusterMax", 36, "Maximum value of chi2 fit over ITS clusters for pi <- charm baryon"};
131139

140+
// ML inference
141+
Configurable<bool> applyMl{"applyMl", true, "Flag to apply ML selections"};
142+
Configurable<std::vector<double>> binsPtMl{"binsPtMl", std::vector<double>{hf_cuts_ml::vecBinsPt}, "pT bin limits for ML application"};
143+
Configurable<std::vector<int>> cutDirMl{"cutDirMl", std::vector<int>{hf_cuts_ml::vecCutDir}, "Whether to reject score values greater or smaller than the threshold"};
144+
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"};
145+
Configurable<int> nClassesMl{"nClassesMl", static_cast<int>(hf_cuts_ml::NCutScores), "Number of classes in ML model"};
146+
Configurable<std::vector<std::string>> namesInputFeatures{"namesInputFeatures", std::vector<std::string>{"feature1", "feature2"}, "Names of ML model input features"};
147+
148+
// CCDB configuration
149+
Configurable<std::string> ccdbUrl{"ccdbUrl", "http://alice-ccdb.cern.ch", "url of the ccdb repository"};
150+
Configurable<std::vector<std::string>> modelPathsCCDB{"modelPathsCCDB", std::vector<std::string>{"EventFiltering/PWGHF/BDTXic0ToXipiKf"}, "Paths of models on CCDB"};
151+
Configurable<std::vector<std::string>> onnxFileNames{"onnxFileNames", std::vector<std::string>{"ModelHandler_onnx_Xic0ToXipiKf.onnx"}, "ONNX file names for each pT bin (if not from CCDB full path)"};
152+
Configurable<int64_t> timestampCCDB{"timestampCCDB", -1, "timestamp of the ONNX file for ML model used to query in CCDB"};
153+
Configurable<bool> loadModelsFromCCDB{"loadModelsFromCCDB", false, "Flag to enable or disable the loading of models from CCDB"};
154+
155+
o2::analysis::HfMlResponseXic0ToXiPi<float> hfMlResponse;
156+
std::vector<float> outputMlXic0ToXiPi = {};
157+
o2::ccdb::CcdbApi ccdbApi;
158+
132159
TrackSelectorPi selectorPion;
133160
TrackSelectorPr selectorProton;
134161

@@ -189,6 +216,18 @@ struct HfCandidateSelectorToXiPi {
189216
registry.add("hSelMassCharmBaryon", "hSelMassCharmBaryon;status;entries", {HistType::kTH1F, {axisSel}});
190217
registry.add("hSelDcaXYToPvV0Daughters", "hSelDcaXYToPvV0Daughters;status;entries", {HistType::kTH1F, {axisSel}});
191218
registry.add("hSelDcaXYToPvPiFromCasc", "hSelDcaXYToPvPiFromCasc;status;entries", {HistType::kTH1F, {axisSel}});
219+
220+
if (applyMl) {
221+
hfMlResponse.configure(binsPtMl, cutsMl, cutDirMl, nClassesMl);
222+
if (loadModelsFromCCDB) {
223+
ccdbApi.init(ccdbUrl);
224+
hfMlResponse.setModelPathsCCDB(onnxFileNames, ccdbApi, modelPathsCCDB, timestampCCDB);
225+
} else {
226+
hfMlResponse.setModelPathsLocal(onnxFileNames);
227+
}
228+
hfMlResponse.cacheInputFeaturesIndices(namesInputFeatures);
229+
hfMlResponse.init();
230+
}
192231
}
193232

194233
void process(aod::HfCandToXiPi const& candidates,
@@ -204,6 +243,7 @@ struct HfCandidateSelectorToXiPi {
204243

205244
bool resultSelections = true; // True if the candidate passes all the selections, False otherwise
206245

246+
auto ptCand = RecoDecay::pt(candidate.pxCharmBaryon(), candidate.pyCharmBaryon());
207247
auto trackV0PosDauId = candidate.posTrackId(); // positive V0 daughter
208248
auto trackV0NegDauId = candidate.negTrackId(); // negative V0 daughter
209249
auto trackPiFromCascId = candidate.bachelorId(); // pion <- cascade
@@ -518,6 +558,17 @@ struct HfCandidateSelectorToXiPi {
518558
} else {
519559
registry.fill(HIST("hSelMassCharmBaryon"), 0);
520560
}
561+
562+
// ML selections
563+
if (applyMl) {
564+
bool isSelectedMlXic0 = false;
565+
std::vector<float> inputFeaturesXic0 = hfMlResponse.getInputFeatures(candidate, trackPiFromLam, trackPiFromCasc, trackPiFromCharm);
566+
isSelectedMlXic0 = hfMlResponse.isSelectedMl(inputFeaturesXic0, ptCand, outputMlXic0ToXiPi);
567+
if (!isSelectedMlXic0) {
568+
continue;
569+
}
570+
hfMlToXiPi(outputMlXic0ToXiPi);
571+
}
521572

522573
hfSelToXiPi(statusPidLambda, statusPidCascade, statusPidCharmBaryon, statusInvMassLambda, statusInvMassCascade, statusInvMassCharmBaryon, resultSelections, infoTpcStored, infoTofStored,
523574
trackPiFromCharm.tpcNSigmaPi(), trackPiFromCasc.tpcNSigmaPi(), trackPiFromLam.tpcNSigmaPi(), trackPrFromLam.tpcNSigmaPr(),

‎PWGHF/TableProducer/treeCreatorToXiPi.cxx‎

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,8 @@ DECLARE_SOA_COLUMN(PzPiFromCharmBaryon, pzPiFromCharmBaryon, float);
7575
DECLARE_SOA_COLUMN(PxLambda, pxLambda, float);
7676
DECLARE_SOA_COLUMN(PyLambda, pyLambda, float);
7777
DECLARE_SOA_COLUMN(PzLambda, pzLambda, float);
78+
DECLARE_SOA_COLUMN(PtCharmBaryon, ptCharmBaryon, float);
79+
DECLARE_SOA_COLUMN(PtPiFromCharmBaryon, ptPiFromCharmBaryon, float);
7880
DECLARE_SOA_COLUMN(PxPiFromCasc, pxPiFromCasc, float);
7981
DECLARE_SOA_COLUMN(PyPiFromCasc, pyPiFromCasc, float);
8082
DECLARE_SOA_COLUMN(PzPiFromCasc, pzPiFromCasc, float);
@@ -169,6 +171,7 @@ DECLARE_SOA_TABLE(HfToXiPiFulls, "AOD", "HFTOXIPIFULL",
169171
full::XDecayVtxV0, full::YDecayVtxV0, full::ZDecayVtxV0,
170172
full::SignDecay,
171173
full::CovVtxCharmBaryonXX, full::CovVtxCharmBaryonYY, full::CovVtxCharmBaryonZZ,
174+
full::PtCharmBaryon, full::PtPiFromCharmBaryon,
172175
full::PxCharmBaryon, full::PyCharmBaryon, full::PzCharmBaryon,
173176
full::PxCasc, full::PyCasc, full::PzCasc,
174177
full::PxPiFromCharmBaryon, full::PyPiFromCharmBaryon, full::PzPiFromCharmBaryon,
@@ -202,15 +205,16 @@ DECLARE_SOA_TABLE(HfToXiPiLites, "AOD", "HFTOXIPILITE",
202205
full::XDecayVtxCascade, full::YDecayVtxCascade, full::ZDecayVtxCascade,
203206
full::XDecayVtxV0, full::YDecayVtxV0, full::ZDecayVtxV0,
204207
full::SignDecay,
205-
full::PxCharmBaryon, full::PyCharmBaryon, full::PzCharmBaryon,
206-
full::PxPiFromCharmBaryon, full::PyPiFromCharmBaryon, full::PzPiFromCharmBaryon,
208+
full::PtCharmBaryon, full::PtPiFromCharmBaryon,
207209
full::PxPiFromCasc, full::PyPiFromCasc, full::PzPiFromCasc,
208210
full::PxPosV0Dau, full::PyPosV0Dau, full::PzPosV0Dau,
209211
full::PxNegV0Dau, full::PyNegV0Dau, full::PzNegV0Dau,
210212
full::ImpactParCascXY, full::ImpactParPiFromCharmBaryonXY,
211213
full::ErrImpactParCascXY, full::ErrImpactParPiFromCharmBaryonXY,
212214
full::InvMassLambda, full::InvMassCascade, full::InvMassCharmBaryon,
215+
full::CosPAV0, full::CosPACharmBaryon, full::CosPACasc,
213216
full::EtaV0PosDau, full::EtaV0NegDau, full::EtaPiFromCasc, full::EtaPiFromCharmBaryon,
217+
full::EtaCharmBaryon,
214218
full::DcaXYToPvV0Dau0, full::DcaXYToPvV0Dau1, full::DcaXYToPvCascDau,
215219
full::DcaCascDau, full::DcaV0Dau, full::DcaCharmBaryonDau,
216220
full::ErrorDecayLengthCharmBaryon, full::NormImpParCascade, full::NormImpParPiFromCharmBar,
@@ -283,6 +287,8 @@ struct HfTreeCreatorToXiPi {
283287
candidate.covVtxCharmBaryon0(),
284288
candidate.covVtxCharmBaryon3(),
285289
candidate.covVtxCharmBaryon5(),
290+
RecoDecay::pt(candidate.pxCharmBaryon(), candidate.pyCharmBaryon()),
291+
RecoDecay::pt(candidate.pxBachFromCharmBaryon(), candidate.pyBachFromCharmBaryon()),
286292
candidate.pxCharmBaryon(),
287293
candidate.pyCharmBaryon(),
288294
candidate.pzCharmBaryon(),
@@ -404,15 +410,11 @@ struct HfTreeCreatorToXiPi {
404410
candidate.yDecayVtxV0(),
405411
candidate.zDecayVtxV0(),
406412
candidate.signDecay(),
413+
RecoDecay::pt(candidate.pxCharmBaryon(), candidate.pyCharmBaryon()),
414+
RecoDecay::pt(candidate.pxBachFromCharmBaryon(), candidate.pyBachFromCharmBaryon()),
407415
candidate.pxCharmBaryon(),
408416
candidate.pyCharmBaryon(),
409417
candidate.pzCharmBaryon(),
410-
candidate.pxBachFromCharmBaryon(),
411-
candidate.pyBachFromCharmBaryon(),
412-
candidate.pzBachFromCharmBaryon(),
413-
candidate.pxBachFromCasc(),
414-
candidate.pyBachFromCasc(),
415-
candidate.pzBachFromCasc(),
416418
candidate.pxPosV0Dau(),
417419
candidate.pyPosV0Dau(),
418420
candidate.pzPosV0Dau(),
@@ -426,10 +428,14 @@ struct HfTreeCreatorToXiPi {
426428
candidate.invMassLambda(),
427429
candidate.invMassCascade(),
428430
candidate.invMassCharmBaryon(),
431+
candidate.cosPAV0(),
432+
candidate.cosPACharmBaryon(),
433+
candidate.cosPACasc(),
429434
candidate.etaV0PosDau(),
430435
candidate.etaV0NegDau(),
431436
candidate.etaBachFromCasc(),
432437
candidate.etaBachFromCharmBaryon(),
438+
candidate.etaCharmBaryon(),
433439
candidate.dcaXYToPvV0Dau0(),
434440
candidate.dcaXYToPvV0Dau1(),
435441
candidate.dcaXYToPvCascDau(),

0 commit comments

Comments
 (0)