Skip to content

Commit 5c54152

Browse files
authored
DPG/AOTTrack: add configurables for the ML cut application (#6979)
1 parent 66c0be2 commit 5c54152

1 file changed

Lines changed: 3 additions & 2 deletions

File tree

DPG/Tasks/AOTTrack/tagAndProbeDmesons.cxx

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -226,6 +226,8 @@ struct TagTwoProngDisplacedVertices {
226226
Configurable<std::vector<std::string>> onnxFileNamesDzeroFromDstar{"onnxFileNamesDzeroFromDstar", std::vector<std::string>{"ModelHandler_onnx_DzeroToKPi.onnx"}, "ONNX file names for Kpi pairs from D0 <- D*+ decays"};
227227
Configurable<std::vector<std::string>> onnxFileNamesDzeroKaKaFromDstar{"onnxFileNamesDzeroKaKaFromDstar", std::vector<std::string>{"ModelHandler_onnx_DzeroToKK.onnx"}, "ONNX file names for KK pairs from D0 <- D*+ decays"};
228228
// ML cuts
229+
Configurable<int> numMlClasses{"numMlClasses", 3, "Number of classes for the ML models"};
230+
Configurable<std::vector<int>> cutDirMl{"cutDirMl", std::vector<int>{o2::cuts_ml::CutDirection::CutGreater, o2::cuts_ml::CutDirection::CutSmaller, o2::cuts_ml::CutDirection::CutSmaller}, "Whether to reject score values greater or smaller than the threshold"};
229231
Configurable<LabeledArray<double>> mlCutsPiPiFromDplus{"mlCutsPiPiFromDplus", {aod::tagandprobe::mlCuts[0], aod::tagandprobe::nBinsPt, 3, aod::tagandprobe::labelsEmpty, aod::tagandprobe::labelsMlScores}, "ML Selections for pipi pairs from D+ decays"};
230232
Configurable<LabeledArray<double>> mlCutsKaKaFromDsOrDplus{"mlCutsKaKaFromDsOrDplus", {aod::tagandprobe::mlCuts[0], aod::tagandprobe::nBinsPt, 3, aod::tagandprobe::labelsEmpty, aod::tagandprobe::labelsMlScores}, "ML Selections for KK pairs from Ds or D+ decays"};
231233
Configurable<LabeledArray<double>> mlCutsDzeroFromDstar{"mlCutsDzeroFromDstar", {aod::tagandprobe::mlCuts[0], aod::tagandprobe::nBinsPt, 3, aod::tagandprobe::labelsEmpty, aod::tagandprobe::labelsMlScores}, "ML Selections for Kpi pairs from D0 <- D*+ decays"};
@@ -320,14 +322,13 @@ struct TagTwoProngDisplacedVertices {
320322
registry.add<TH2>("hMassDzeroKaKaVsPt", ";#it{p}_{T}(K#pi) (GeV/#it{c}); #it{M}(K#pi) (GeV/#it{c}^{2})", HistType::kTH2D, {axisPt, axisMassKaPi});
321323
}
322324

323-
const std::vector<int> cutDirMl = {o2::cuts_ml::CutDirection::CutGreater, o2::cuts_ml::CutDirection::CutSmaller, o2::cuts_ml::CutDirection::CutSmaller};
324325
const std::array<LabeledArray<double>, aod::tagandprobe::TagChannels::NTagChannels> mlCuts = {mlConfig.mlCutsPiPiFromDplus, mlConfig.mlCutsKaKaFromDsOrDplus, mlConfig.mlCutsDzeroFromDstar, mlConfig.mlCutsDzeroFromDstar, mlConfig.mlCutsDzeroKaKaFromDstar};
325326
const std::array<std::vector<std::string>, aod::tagandprobe::TagChannels::NTagChannels> onnxFileNames = {mlConfig.onnxFileNamesPiPiFromDplus, mlConfig.onnxFileNamesKaKaFromDsOrDplus, mlConfig.onnxFileNamesDzeroFromDstar, mlConfig.onnxFileNamesDzeroFromDstar, mlConfig.onnxFileNamesDzeroKaKaFromDstar};
326327
const std::array<std::vector<std::string>, aod::tagandprobe::TagChannels::NTagChannels> modelPathsCcdb = {mlConfig.modelPathsCcdbPiPiFromDplus, mlConfig.modelPathsCcdbKaKaFromDsOrDplus, mlConfig.modelPathsCcdbDzeroFromDstar, mlConfig.modelPathsCcdbDzeroFromDstar, mlConfig.modelPathsCcdbDzeroKaKaFromDstar};
327328
applyMl = {mlConfig.applyMlPiPiFromDplus, mlConfig.applyMlKaKaFromDsOrDplus, mlConfig.applyMlDzeroFromDstar, mlConfig.applyMlDzeroFromDstar, mlConfig.applyMlDzeroKaKaFromDstar};
328329
for (auto iChannel{0u}; iChannel < aod::tagandprobe::TagChannels::NTagChannels; ++iChannel) {
329330
if (applyMl[iChannel]) {
330-
mlResponse[iChannel].configure(ptBinsForTopologicalCuts[iChannel], mlCuts[iChannel], cutDirMl, 3);
331+
mlResponse[iChannel].configure(ptBinsForTopologicalCuts[iChannel], mlCuts[iChannel], mlConfig.cutDirMl, mlConfig.numMlClasses);
331332
if (mlConfig.loadMlModelsFromCcdb) {
332333
mlResponse[iChannel].setModelPathsCCDB(onnxFileNames[iChannel], ccdbApi, modelPathsCcdb[iChannel], mlConfig.timestampCcdbForMlModels);
333334
} else {

0 commit comments

Comments
 (0)