diff --git a/PWGLF/DataModel/LFStrangenessMLTables.h b/PWGLF/DataModel/LFStrangenessMLTables.h index ad30c86524b..7784a521dd7 100644 --- a/PWGLF/DataModel/LFStrangenessMLTables.h +++ b/PWGLF/DataModel/LFStrangenessMLTables.h @@ -16,6 +16,7 @@ #include #include +#include // Creating output TTree for ML analysis namespace o2::aod @@ -237,8 +238,8 @@ DECLARE_SOA_TABLE(CascMLCandidates, "AOD", "CAMLCANDIDATES", namespace CascMLSelection { -DECLARE_SOA_COLUMN(XiBDTScore, xiBDTScore, float); -DECLARE_SOA_COLUMN(OmegaBDTScore, omegaBDTScore, float); +DECLARE_SOA_COLUMN(XiBDTScore, xiBDTScore, std::vector); +DECLARE_SOA_COLUMN(OmegaBDTScore, omegaBDTScore, std::vector); } // namespace CascMLSelection DECLARE_SOA_TABLE(CascXiMLScores, "AOD", "CASCXIMLSCORES", diff --git a/PWGLF/TableProducer/Strangeness/cascademlselection.cxx b/PWGLF/TableProducer/Strangeness/cascademlselection.cxx index dd2341964f0..9f096144fb2 100644 --- a/PWGLF/TableProducer/Strangeness/cascademlselection.cxx +++ b/PWGLF/TableProducer/Strangeness/cascademlselection.cxx @@ -9,21 +9,18 @@ // granted to it by virtue of its status as an Intergovernmental Organization // or submit itself to any jurisdiction. // -// *+-+*+-+*+-+*+-+*+-+*+-+*+-+*+-+* -// Lambdakzero ML selection task -// *+-+*+-+*+-+*+-+*+-+*+-+*+-+*+-+* -// -// Comments, questions, complaints, suggestions? -// Please write to: -// gianni.shigeru.setoue.liveraro@cern.ch -// romain.schotter@cern.ch -// david.dobrigkeit.chinellato@cern.ch +/// \file Cascade ML selection task +/// \brief Produces ML response table for cascade selection at analysis level, either when running over original data or derived data. +/// \author Gianni Shigeru Setoue Liveraro Catalano , UNICAMP +/// \author Romain Schotter , Austrian Academy of Sciences +/// \author David Dobrigkeit Chinellato , Austrian Academy of Sciences // #include "PWGLF/DataModel/LFStrangenessMLTables.h" #include "PWGLF/DataModel/LFStrangenessTables.h" +#include "PWGLF/Utils/CascadeMlResponse.h" -#include "Tools/ML/model.h" +#include "Tools/ML/MlResponse.h" #include #include @@ -31,6 +28,7 @@ #include #include #include +#include #include #include #include @@ -38,10 +36,16 @@ #include #include +#include +#include + +#include #include #include #include +#include #include +#include #include #include @@ -51,34 +55,39 @@ using namespace o2::framework::expressions; using namespace o2::ml; // For original data loops -using CascOriginalDatas = soa::Join; +using CascOriginalDatas = soa::Join; // For derived data analysis -using CascDerivedDatas = soa::Join; +using CascDerivedDatas = soa::Join; struct cascademlselection { - o2::ml::OnnxModel mlModelXiMinus; - o2::ml::OnnxModel mlModelXiPlus; - o2::ml::OnnxModel mlModelOmegaMinus; - o2::ml::OnnxModel mlModelOmegaPlus; + o2::analysis::CascadeMlResponse mlModelXiMinus; + o2::analysis::CascadeMlResponse mlModelXiPlus; + o2::analysis::CascadeMlResponse mlModelOmegaMinus; + o2::analysis::CascadeMlResponse mlModelOmegaPlus; // Custom grouping std::vector> cascadesGrouped; - std::map metadata; - Produces xiMLSelections; // optionally aggregate information from ML output for posterior analysis (derived data) Produces omegaMLSelections; // optionally aggregate information from ML output for posterior analysis (derived data) HistogramRegistry histos{"Histos", {}, OutputObjHandlingPolicy::AnalysisObject}; + // BDT score histograms, indexed [pT bin][class], one set per particle species + std::vector>> histScoreXiMinus; + std::vector>> histScoreXiPlus; + std::vector>> histScoreOmegaMinus; + std::vector>> histScoreOmegaPlus; + // CCDB configuration o2::ccdb::CcdbApi ccdbApi; Service ccdb; - int mRunNumber; + int mRunNumber = -1; // CCDB options struct : ConfigurableGroup { + std::string prefix = "ccdbConfigurations"; Configurable ccdburl{"ccdb-url", "http://alice-ccdb.cern.ch", "url of the ccdb repository"}; Configurable grpPath{"grpPath", "GLO/GRP/GRP", "Path of the grp file"}; Configurable grpmagPath{"grpmagPath", "GLO/Config/GRPMagField", "CCDB path of the GRPMagField object"}; @@ -88,29 +97,43 @@ struct cascademlselection { // Machine learning evaluation for pre-selection and corresponding information generation struct : ConfigurableGroup { + std::string prefix = "mlConfigurations"; // ML classifiers: master flags to populate ML Selection tables Configurable calculateXiMinusScores{"mlConfigurations.calculateXiMinusScores", true, "calculate XiMinus ML scores"}; Configurable calculateXiPlusScores{"mlConfigurations.calculateXiPlusScores", true, "calculate XiPlus ML scores"}; Configurable calculateOmegaMinusScores{"mlConfigurations.calculateOmegaMinusScores", true, "calculate OmegaMinus ML scores"}; Configurable calculateOmegaPlusScores{"mlConfigurations.calculateOmegaPlusScores", true, "calculate OmegaPlus ML scores"}; - // ML input for ML calculation - Configurable modelPathCCDB{"mlConfigurations.modelPathCCDB", "", "ML Model path in CCDB"}; + // List and order of input features fed to the ONNX models; any subset/order of the names + // registered in CascMlResponse::setAvailableInputFeatures can be used here + Configurable> namesInputFeatures{"mlConfigurations.namesInputFeatures", std::vector{"cascradius", "v0radius", "casccosPA", "v0cosPA", "dcapostopv", "dcanegtopv", "dcabachtopv", "dcacascdaughters", "dcaV0daughters", "dcav0topv", "bachBaryonCosPA", "bachBaryonDCAxyToPV"}, "Names (and order) of the input features to be used in the ML models"}; + + // ML input for ML calculation: one model (ONNX file / CCDB path) per pT bin + Configurable> modelPathsCCDBXiMinus{"mlConfigurations.modelPathsCCDBXiMinus", std::vector{""}, "ML Model paths in CCDB for Xi-. One per pT bin."}; + Configurable> modelPathsCCDBXiPlus{"mlConfigurations.modelPathsCCDBXiPlus", std::vector{""}, "ML Model paths in CCDB for Xi+. One per pT bin."}; + Configurable> modelPathsCCDBOmegaMinus{"mlConfigurations.modelPathsCCDBOmegaMinus", std::vector{""}, "ML Model paths in CCDB for Omega-. One per pT bin."}; + Configurable> modelPathsCCDBOmegaPlus{"mlConfigurations.modelPathsCCDBOmegaPlus", std::vector{""}, "ML Model paths in CCDB for Omega+. One per pT bin."}; Configurable timestampCCDB{"mlConfigurations.timestampCCDB", -1, "timestamp of the ONNX file for ML model used to query in CCDB. Exceptions: > 0 for the specific timestamp, 0 gets the run dependent timestamp"}; Configurable loadModelsFromCCDB{"mlConfigurations.loadModelsFromCCDB", false, "Flag to enable or disable the loading of models from CCDB"}; Configurable enableOptimizations{"mlConfigurations.enableOptimizations", false, "Enables the ONNX extended model-optimization: sessionOptions.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED)"}; - // Local paths for test purposes - Configurable localModelPathXiMinus{"mlConfigurations.localModelPathXiMinus", "XiMinus_BDTModel.onnx", "(std::string) Path to the local .onnx file."}; - Configurable localModelPathXiPlus{"mlConfigurations.localModelPathXiPlus", "XiPlus_BDTModel.onnx", "(std::string) Path to the local .onnx file."}; - Configurable localModelPathOmegaMinus{"mlConfigurations.localModelPathOmegaMinus", "OmegaMinus_BDTModel.onnx", "(std::string) Path to the local .onnx file."}; - Configurable localModelPathOmegaPlus{"mlConfigurations.localModelPathOmegaPlus", "OmegaPlus_BDTModel.onnx", "(std::string) Path to the local .onnx file."}; - - // Thresholds for choosing to populate V0Cores tables with pre-selections - Configurable thresholdXiMinus{"mlConfigurations.thresholdXiMinus", -1.0f, "Threshold to keep XiMinus candidates"}; - Configurable thresholdXiPlus{"mlConfigurations.thresholdXiPlus", -1.0f, "Threshold to keep XiPlus candidates"}; - Configurable thresholdOmegaMinus{"mlConfigurations.thresholdOmegaMinus", -1.0f, "Threshold to keep OmegaMinus candidates"}; - Configurable thresholdOmegaPlus{"mlConfigurations.thresholdOmegaPlus", -1.0f, "Threshold to keep OmegaPlus candidates"}; + // Local/cvmfs paths (also used as CCDB download destination filenames), one per pT bin + Configurable> onnxFileNamesXiMinus{"mlConfigurations.onnxFileNamesXiMinus", std::vector{"XiMinus_BDTModel.onnx"}, "(std::string) Paths to the local .onnx file. One per pT bin."}; + Configurable> onnxFileNamesXiPlus{"mlConfigurations.onnxFileNamesXiPlus", std::vector{"XiPlus_BDTModel.onnx"}, "(std::string) Paths to the local .onnx file. One per pT bin."}; + Configurable> onnxFileNamesOmegaMinus{"mlConfigurations.onnxFileNamesOmegaMinus", std::vector{"OmegaMinus_BDTModel.onnx"}, "(std::string) Paths to the local .onnx file. One per pT bin."}; + Configurable> onnxFileNamesOmegaPlus{"mlConfigurations.onnxFileNamesOmegaPlus", std::vector{"OmegaPlus_BDTModel.onnx"}, "(std::string) Paths to the local .onnx file. One per pT bin."}; + + // Binning + Configurable> binsPtXiMinus{"binsPtXiMinus", std::vector{0., 10.}, "pT bin limits for ML application for Xi-"}; + Configurable> binsPtXiPlus{"binsPtXiPlus", std::vector{0., 10.}, "pT bin limits for ML application for Xi+"}; + Configurable> binsPtOmegaMinus{"binsPtOmegaMinus", std::vector{0., 10.}, "pT bin limits for ML application for Omega-"}; + Configurable> binsPtOmegaPlus{"binsPtOmegaPlus", std::vector{0., 10.}, "pT bin limits for ML application for Omega+"}; + + // Number of classes in the ML models. Default is 2 (signal and background) + Configurable nClassesMlXiMinus{"nClassesMlXiMinus", 2, "Number of classes in ML model for Xi-"}; + Configurable nClassesMlXiPlus{"nClassesMlXiPlus", 2, "Number of classes in ML model for Xi+"}; + Configurable nClassesMlOmegaMinus{"nClassesMlOmegaMinus", 2, "Number of classes in ML model for Omega-"}; + Configurable nClassesMlOmegaPlus{"nClassesMlOmegaPlus", 2, "Number of classes in ML model for Omega+"}; } mlConfigurations; // Axis @@ -146,62 +169,35 @@ struct cascademlselection { mlConfigurations.calculateOmegaPlusScores) { if (mlConfigurations.timestampCCDB.value != -1) timeStampML = mlConfigurations.timestampCCDB.value; - LoadMachines(timeStampML); + loadMachines(timeStampML); } } // function to load models for ML-based classifiers - void LoadMachines(int64_t timeStampML) + void loadMachines(int64_t timeStampML) { + auto loadModel = [&](bool doCalculate, o2::analysis::CascadeMlResponse& model, std::vector const& onnxFileNames, std::vector const& pathsCCDB) { + if (!doCalculate) { + return; + } + if (mlConfigurations.loadModelsFromCCDB) { + model.setModelPathsCCDB(onnxFileNames, ccdbApi, pathsCCDB, timeStampML); + } else { + model.setModelPathsLocal(onnxFileNames); + } + model.init(mlConfigurations.enableOptimizations.value); + }; + if (mlConfigurations.loadModelsFromCCDB) { ccdbApi.init(ccdbConfigurations.ccdburl); LOG(info) << "Fetching cascade models for timestamp: " << timeStampML; + } - if (mlConfigurations.calculateXiMinusScores) { - bool retrieveSuccess = ccdbApi.retrieveBlob(mlConfigurations.modelPathCCDB, ".", metadata, timeStampML, false, mlConfigurations.localModelPathXiMinus.value); - if (retrieveSuccess) { - mlModelXiMinus.initModel(mlConfigurations.localModelPathXiMinus.value, mlConfigurations.enableOptimizations.value); - } else { - LOG(fatal) << "Error encountered while fetching/loading the XiMinus model from CCDB! Maybe the model doesn't exist yet for this runnumber/timestamp?"; - } - } - - if (mlConfigurations.calculateXiPlusScores) { - bool retrieveSuccess = ccdbApi.retrieveBlob(mlConfigurations.modelPathCCDB, ".", metadata, timeStampML, false, mlConfigurations.localModelPathXiPlus.value); - if (retrieveSuccess) { - mlModelXiPlus.initModel(mlConfigurations.localModelPathXiPlus.value, mlConfigurations.enableOptimizations.value); - } else { - LOG(fatal) << "Error encountered while fetching/loading the XiPlus model from CCDB! Maybe the model doesn't exist yet for this runnumber/timestamp?"; - } - } - - if (mlConfigurations.calculateOmegaMinusScores) { - bool retrieveSuccess = ccdbApi.retrieveBlob(mlConfigurations.modelPathCCDB, ".", metadata, timeStampML, false, mlConfigurations.localModelPathOmegaMinus.value); - if (retrieveSuccess) { - mlModelOmegaMinus.initModel(mlConfigurations.localModelPathOmegaMinus.value, mlConfigurations.enableOptimizations.value); - } else { - LOG(fatal) << "Error encountered while fetching/loading the OmegaMinus model from CCDB! Maybe the model doesn't exist yet for this runnumber/timestamp?"; - } - } + loadModel(mlConfigurations.calculateXiMinusScores, mlModelXiMinus, mlConfigurations.onnxFileNamesXiMinus, mlConfigurations.modelPathsCCDBXiMinus); + loadModel(mlConfigurations.calculateXiPlusScores, mlModelXiPlus, mlConfigurations.onnxFileNamesXiPlus, mlConfigurations.modelPathsCCDBXiPlus); + loadModel(mlConfigurations.calculateOmegaMinusScores, mlModelOmegaMinus, mlConfigurations.onnxFileNamesOmegaMinus, mlConfigurations.modelPathsCCDBOmegaMinus); + loadModel(mlConfigurations.calculateOmegaPlusScores, mlModelOmegaPlus, mlConfigurations.onnxFileNamesOmegaPlus, mlConfigurations.modelPathsCCDBOmegaPlus); - if (mlConfigurations.calculateOmegaPlusScores) { - bool retrieveSuccess = ccdbApi.retrieveBlob(mlConfigurations.modelPathCCDB, ".", metadata, timeStampML, false, mlConfigurations.localModelPathOmegaPlus.value); - if (retrieveSuccess) { - mlModelOmegaPlus.initModel(mlConfigurations.localModelPathOmegaPlus.value, mlConfigurations.enableOptimizations.value); - } else { - LOG(fatal) << "Error encountered while fetching/loading the OmegaPlus model from CCDB! Maybe the model doesn't exist yet for this runnumber/timestamp?"; - } - } - } else { - if (mlConfigurations.calculateXiMinusScores) - mlModelXiMinus.initModel(mlConfigurations.localModelPathXiMinus.value, mlConfigurations.enableOptimizations.value); - if (mlConfigurations.calculateXiPlusScores) - mlModelXiPlus.initModel(mlConfigurations.localModelPathXiPlus.value, mlConfigurations.enableOptimizations.value); - if (mlConfigurations.calculateOmegaMinusScores) - mlModelOmegaMinus.initModel(mlConfigurations.localModelPathOmegaMinus.value, mlConfigurations.enableOptimizations.value); - if (mlConfigurations.calculateOmegaPlusScores) - mlModelOmegaPlus.initModel(mlConfigurations.localModelPathOmegaPlus.value, mlConfigurations.enableOptimizations.value); - } LOG(info) << "Cascade ML Models loaded."; } @@ -211,45 +207,96 @@ struct cascademlselection { histos.add("hEventVertexZ", "hEventVertexZ", kTH1F, {vertexZ}); ccdb->setURL(ccdbConfigurations.ccdburl); + + // builds a shape-valid but functionally unused cuts array: this task only stores + // raw ML scores (thresholds are applied downstream), so no cut direction is used + auto dummyCuts = [](int nBins, int nClasses) { + std::vector zeros(static_cast(nBins) * static_cast(nClasses), 0.); + return LabeledArray(zeros.data(), nBins, nClasses); + }; + + auto configureModel = [&](o2::analysis::CascadeMlResponse& model, std::vector const& binsPt, int nClasses) { + int nBins = static_cast(binsPt.size()) - 1; + model.configure(binsPt, dummyCuts(nBins, nClasses), std::vector(nClasses, o2::cuts_ml::CutNot), nClasses); + model.cacheInputFeaturesIndices(mlConfigurations.namesInputFeatures); + }; + + auto bookScoreHistos = [&](std::vector>>& target, std::vector const& binsPt, int nClasses, std::string const& particleName) { + int nBins = static_cast(binsPt.size()) - 1; + target.resize(nBins); + for (int iBin = 0; iBin < nBins; iBin++) { + target[iBin].resize(nClasses); + for (int iClass = 0; iClass < nClasses; iClass++) { + target[iBin][iClass] = histos.add(Form("BDTScore/%s/pTbin%d/class%d", particleName.c_str(), iBin, iClass), + Form("%s BDT score, %.2f #leq p_{T} < %.2f GeV/c, class %d;BDT score;entries", particleName.c_str(), binsPt[iBin], binsPt[iBin + 1], iClass), + kTH1F, {{100, 0., 1.}}); + } + } + }; + + if (mlConfigurations.calculateXiMinusScores) { + configureModel(mlModelXiMinus, mlConfigurations.binsPtXiMinus, mlConfigurations.nClassesMlXiMinus); + bookScoreHistos(histScoreXiMinus, mlConfigurations.binsPtXiMinus, mlConfigurations.nClassesMlXiMinus, "XiMinus"); + } + if (mlConfigurations.calculateXiPlusScores) { + configureModel(mlModelXiPlus, mlConfigurations.binsPtXiPlus, mlConfigurations.nClassesMlXiPlus); + bookScoreHistos(histScoreXiPlus, mlConfigurations.binsPtXiPlus, mlConfigurations.nClassesMlXiPlus, "XiPlus"); + } + if (mlConfigurations.calculateOmegaMinusScores) { + configureModel(mlModelOmegaMinus, mlConfigurations.binsPtOmegaMinus, mlConfigurations.nClassesMlOmegaMinus); + bookScoreHistos(histScoreOmegaMinus, mlConfigurations.binsPtOmegaMinus, mlConfigurations.nClassesMlOmegaMinus, "OmegaMinus"); + } + if (mlConfigurations.calculateOmegaPlusScores) { + configureModel(mlModelOmegaPlus, mlConfigurations.binsPtOmegaPlus, mlConfigurations.nClassesMlOmegaPlus); + bookScoreHistos(histScoreOmegaPlus, mlConfigurations.binsPtOmegaPlus, mlConfigurations.nClassesMlOmegaPlus, "OmegaPlus"); + } } - // Process candidate and store properties in object - template - void processCandidate(TCascObject const& cand) + // Finds the pT bin matching the MlResponse internal convention (upper_bound over bin edges); + // returns -1 if pt falls outside the configured range, since evaluating the model in that + // case would otherwise trigger a LOG(fatal) inside MlResponse::getModelOutput + int findPtBin(std::vector const& binsPt, float pt) { - // Select features - // FIXME THIS NEEDS ADJUSTING - std::vector inputFeatures{0.0f, 0.0f, - 0.0f, 0.0f}; - - // calculate scores - if (cand.sign() < 0) { - if (mlConfigurations.calculateXiMinusScores) { - float* xiMinusProbability = mlModelXiMinus.evalModel(inputFeatures); - xiMLSelections(xiMinusProbability[1]); - } else { - xiMLSelections(-1); - } - if (mlConfigurations.calculateOmegaMinusScores) { - float* omegaMinusProbability = mlModelOmegaMinus.evalModel(inputFeatures); - omegaMLSelections(omegaMinusProbability[1]); - } else { - omegaMLSelections(-1); - } + if (pt < binsPt.front() || pt >= binsPt.back()) { + return -1; } - if (cand.sign() > 0) { - if (mlConfigurations.calculateXiPlusScores) { - float* xiPlusProbability = mlModelXiPlus.evalModel(inputFeatures); - xiMLSelections(xiPlusProbability[1]); - } else { - xiMLSelections(-1); - } - if (mlConfigurations.calculateOmegaPlusScores) { - float* omegaPlusProbability = mlModelOmegaPlus.evalModel(inputFeatures); - omegaMLSelections(omegaPlusProbability[1]); - } else { - omegaMLSelections(-1); - } + return static_cast(std::distance(binsPt.begin(), std::upper_bound(binsPt.begin(), binsPt.end(), pt))) - 1; + } + + // Evaluates one particle hypothesis' model for a candidate, fills the per-bin/per-class + // score histograms, and returns the signal (class 1) score to be stored in the output table + template + std::vector evaluateModel(o2::analysis::CascadeMlResponse& model, std::vector>>& scoreHistos, std::vector const& binsPt, bool doCalculate, TCascObject const& casc, float pt, TCollision const& coll) + { + if (!doCalculate) { + return {}; + } + int iBin = findPtBin(binsPt, pt); + if (iBin < 0) { + return {}; + } + + auto inputFeatures = model.getInputFeatures(casc, coll); + std::vector output; + model.isSelectedMl(inputFeatures, pt, output); + + for (std::size_t iClass = 0; iClass < output.size() && iClass < scoreHistos[iBin].size(); iClass++) { + scoreHistos[iBin][iClass]->Fill(output[iClass]); + } + return output; + } + + // Process candidate and store properties in object + template + void processCandidate(TCascObject const& casc, float pt, TCollision const& coll) + { + // calculate scores (cascades only ever carry sign +1 or -1) + if (casc.sign() < 0) { + xiMLSelections(evaluateModel(mlModelXiMinus, histScoreXiMinus, mlConfigurations.binsPtXiMinus, mlConfigurations.calculateXiMinusScores, casc, pt, coll)); + omegaMLSelections(evaluateModel(mlModelOmegaMinus, histScoreOmegaMinus, mlConfigurations.binsPtOmegaMinus, mlConfigurations.calculateOmegaMinusScores, casc, pt, coll)); + } else { + xiMLSelections(evaluateModel(mlModelXiPlus, histScoreXiPlus, mlConfigurations.binsPtXiPlus, mlConfigurations.calculateXiPlusScores, casc, pt, coll)); + omegaMLSelections(evaluateModel(mlModelOmegaPlus, histScoreOmegaPlus, mlConfigurations.binsPtOmegaPlus, mlConfigurations.calculateOmegaPlusScores, casc, pt, coll)); } } @@ -269,11 +316,7 @@ struct cascademlselection { histos.fill(HIST("hEventVertexZ"), collision.posZ()); for (std::size_t i = 0; i < cascadesGrouped[collision.globalIndex()].size(); i++) { auto casc = cascades.rawIteratorAt(cascadesGrouped[collision.globalIndex()][i]); - nCandidates++; - if (nCandidates % 50000 == 0) { - LOG(info) << "Candidates processed: " << nCandidates; - } - processCandidate(casc); + processCandidate(casc, casc.pt(), collision); } } } @@ -293,11 +336,7 @@ struct cascademlselection { histos.fill(HIST("hEventVertexZ"), collision.posZ()); for (std::size_t i = 0; i < cascadesGrouped[collision.globalIndex()].size(); i++) { auto casc = cascades.rawIteratorAt(cascadesGrouped[collision.globalIndex()][i]); - nCandidates++; - if (nCandidates % 50000 == 0) { - LOG(info) << "Candidates processed: " << nCandidates; - } - processCandidate(casc); + processCandidate(casc, casc.pt(), collision); } } } diff --git a/PWGLF/Utils/CascadeMlResponse.h b/PWGLF/Utils/CascadeMlResponse.h new file mode 100644 index 00000000000..ff72cf922f9 --- /dev/null +++ b/PWGLF/Utils/CascadeMlResponse.h @@ -0,0 +1,120 @@ +// Copyright 2019-2020 CERN and copyright holders of ALICE O2. +// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders. +// All rights not expressly granted are reserved. +// +// This software is distributed under the terms of the GNU General Public +// License v3 (GPL Version 3), copied verbatim in the file "COPYING". +// +// In applying this license CERN does not waive the privileges and immunities +// granted to it by virtue of its status as an Intergovernmental Organization +// or submit itself to any jurisdiction. + +/// \file CascadeMlResponse.h +/// \brief Class to compute the ML response for Strangeness-analysis selections +/// \author Gianni Shigeru Setoue Liveraro Catalano , UNICAMP +/// \author Romain Schotter , Austrian Academy of Sciences +/// \author David Dobrigkeit Chinellato , Austrian Academy of Sciences + +#ifndef PWGLF_UTILS_CASCADEMLRESPONSE_H_ +#define PWGLF_UTILS_CASCADEMLRESPONSE_H_ + +#include "Tools/ML/MlResponse.h" + +#include + +namespace o2::analysis +{ +// list of input features that can be requested via the mlConfigurations.namesInputFeatures configurable +enum class InputFeaturesCasc : uint8_t { + cascradius = 0, + v0radius, + casccosPA, + v0cosPA, + dcapostopv, + dcanegtopv, + dcabachtopv, + dcacascdaughters, + dcaV0daughters, + dcav0topv, + bachBaryonCosPA, + bachBaryonDCAxyToPV +}; + +template +class CascadeMlResponse : public MlResponse +{ + public: + CascadeMlResponse() = default; + ~CascadeMlResponse() override = default; + + template + std::vector getInputFeatures(TCascade const& casc, TCollision const& coll) + { + std::vector inputFeatures; + inputFeatures.reserve(MlResponse::mCachedIndices.size()); + + for (const auto& idx : MlResponse::mCachedIndices) { + switch (static_cast(idx)) { + case InputFeaturesCasc::cascradius: + inputFeatures.emplace_back(casc.cascradius()); + break; + case InputFeaturesCasc::v0radius: + inputFeatures.emplace_back(casc.v0radius()); + break; + case InputFeaturesCasc::casccosPA: + inputFeatures.emplace_back(casc.casccosPA(coll.posX(), coll.posY(), coll.posZ())); + break; + case InputFeaturesCasc::v0cosPA: + inputFeatures.emplace_back(casc.v0cosPA(coll.posX(), coll.posY(), coll.posZ())); + break; + case InputFeaturesCasc::dcapostopv: + inputFeatures.emplace_back(casc.dcapostopv()); + break; + case InputFeaturesCasc::dcanegtopv: + inputFeatures.emplace_back(casc.dcanegtopv()); + break; + case InputFeaturesCasc::dcabachtopv: + inputFeatures.emplace_back(casc.dcabachtopv()); + break; + case InputFeaturesCasc::dcacascdaughters: + inputFeatures.emplace_back(casc.dcacascdaughters()); + break; + case InputFeaturesCasc::dcaV0daughters: + inputFeatures.emplace_back(casc.dcaV0daughters()); + break; + case InputFeaturesCasc::dcav0topv: + inputFeatures.emplace_back(casc.dcav0topv(coll.posX(), coll.posY(), coll.posZ())); + break; + case InputFeaturesCasc::bachBaryonCosPA: + inputFeatures.emplace_back(casc.bachBaryonCosPA()); + break; + case InputFeaturesCasc::bachBaryonDCAxyToPV: + inputFeatures.emplace_back(casc.bachBaryonDCAxyToPV()); + break; + } + } + return inputFeatures; + } + + protected: + void setAvailableInputFeatures() override + { + MlResponse::mAvailableInputFeatures = { + {"cascradius", static_cast(InputFeaturesCasc::cascradius)}, + {"v0radius", static_cast(InputFeaturesCasc::v0radius)}, + {"casccosPA", static_cast(InputFeaturesCasc::casccosPA)}, + {"v0cosPA", static_cast(InputFeaturesCasc::v0cosPA)}, + {"dcapostopv", static_cast(InputFeaturesCasc::dcapostopv)}, + {"dcanegtopv", static_cast(InputFeaturesCasc::dcanegtopv)}, + {"dcabachtopv", static_cast(InputFeaturesCasc::dcabachtopv)}, + {"dcacascdaughters", static_cast(InputFeaturesCasc::dcacascdaughters)}, + {"dcaV0daughters", static_cast(InputFeaturesCasc::dcaV0daughters)}, + {"dcav0topv", static_cast(InputFeaturesCasc::dcav0topv)}, + {"bachBaryonCosPA", static_cast(InputFeaturesCasc::bachBaryonCosPA)}, + {"bachBaryonDCAxyToPV", static_cast(InputFeaturesCasc::bachBaryonDCAxyToPV)}}; + } +}; + +} // namespace o2::analysis + +#endif // PWGLF_UTILS_CASCADEMLRESPONSE_H_