From 0c0c8e251df0d04937afb63a531981fef39b3079 Mon Sep 17 00:00:00 2001 From: romainschotter Date: Wed, 15 Jan 2025 23:04:51 +0100 Subject: [PATCH 1/4] Add task for a posteriori cascade ML score calculation --- .../TableProducer/Strangeness/CMakeLists.txt | 5 + .../Strangeness/cascademlselection.cxx | 286 ++++++++++++++++++ 2 files changed, 291 insertions(+) create mode 100644 PWGLF/TableProducer/Strangeness/cascademlselection.cxx diff --git a/PWGLF/TableProducer/Strangeness/CMakeLists.txt b/PWGLF/TableProducer/Strangeness/CMakeLists.txt index aee914fcfff..f7759257f73 100644 --- a/PWGLF/TableProducer/Strangeness/CMakeLists.txt +++ b/PWGLF/TableProducer/Strangeness/CMakeLists.txt @@ -138,6 +138,11 @@ o2physics_add_dpl_workflow(lambdakzeromlselection PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2Physics::MLCore COMPONENT_NAME Analysis) +o2physics_add_dpl_workflow(cascademlselection + SOURCES cascademlselection.cxx + PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2Physics::MLCore + COMPONENT_NAME Analysis) + o2physics_add_dpl_workflow(sigma0builder SOURCES sigma0builder.cxx PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2Physics::MLCore diff --git a/PWGLF/TableProducer/Strangeness/cascademlselection.cxx b/PWGLF/TableProducer/Strangeness/cascademlselection.cxx new file mode 100644 index 00000000000..2add7529ce5 --- /dev/null +++ b/PWGLF/TableProducer/Strangeness/cascademlselection.cxx @@ -0,0 +1,286 @@ +// 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. +// +// *+-+*+-+*+-+*+-+*+-+*+-+*+-+*+-+* +// 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 +// + +#include +#include +#include +#include + +#include "Framework/runDataProcessing.h" +#include "Framework/AnalysisTask.h" +#include "Framework/HistogramRegistry.h" +#include "Framework/AnalysisDataModel.h" +#include "Framework/ASoAHelpers.h" +#include "Framework/ASoA.h" +#include "ReconstructionDataFormats/Track.h" +#include "Common/Core/RecoDecay.h" +#include "Common/Core/trackUtilities.h" +#include "PWGLF/DataModel/LFStrangenessTables.h" +#include "PWGLF/DataModel/LFStrangenessPIDTables.h" +#include "PWGLF/DataModel/LFStrangenessMLTables.h" +#include "Common/Core/TrackSelection.h" +#include "Common/DataModel/TrackSelectionTables.h" +#include "Common/DataModel/EventSelection.h" +#include "Common/DataModel/Centrality.h" +#include "Common/DataModel/PIDResponse.h" +#include "CCDB/BasicCCDBManager.h" +#include +#include +#include +#include +#include +#include +#include "Tools/ML/MlResponse.h" +#include "Tools/ML/model.h" + +using namespace o2; +using namespace o2::analysis; +using namespace o2::framework; +using namespace o2::framework::expressions; +using namespace o2::ml; +using std::array; +using std::cout; +using std::endl; + +// For original data loops +using CascOriginalDatas = soa::Join; + +// For derived data analysis +using CascDerivedDatas = soa::Join; + +struct cascademlselecttion { + o2::ml::OnnxModel mlModelXiMinus; + o2::ml::OnnxModel mlModelXiPlus; + o2::ml::OnnxModel mlModelOmegaMinus; + o2::ml::OnnxModel mlModelOmegaPlus; + + 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}; + + // CCDB configuration + o2::ccdb::CcdbApi ccdbApi; + Service ccdb; + int mRunNumber; + + // CCDB options + struct : ConfigurableGroup { + 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"}; + Configurable lutPath{"lutPath", "GLO/Param/MatLUT", "Path of the Lut parametrization"}; + Configurable geoPath{"geoPath", "GLO/Config/GeometryAligned", "Path of the geometry file"}; + } ccdbConfigurations; + + // Machine learning evaluation for pre-selection and corresponding information generation + struct : ConfigurableGroup { + // ML classifiers: master flags to populate ML Selection tables + Configurable calculateXiMinusScores{"mlConfigurations.calculateXiMinusScores", false, "calculate XiMinus ML scores"}; + Configurable calculateXiPlusScores{"mlConfigurations.calculateXiPlusScores", false, "calculate XiPlus ML scores"}; + Configurable calculateOmegaMinusScores{"mlConfigurations.calculateOmegaMinusScores", false, "calculate OmegaMinus ML scores"}; + Configurable calculateOmegaPlusScores{"mlConfigurations.calculateOmegaPlusScores", false, "calculate OmegaPlus ML scores"}; + + // ML input for ML calculation + Configurable modelPathCCDB{"mlConfigurations.modelPathCCDB", "", "ML Model path in CCDB"}; + 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"}; + } mlConfigurations; + + // Axis + // base properties + ConfigurableAxis vertexZ{"vertexZ", {30, -15.0f, 15.0f}, ""}; + + int nCandidates = 0; + + template + void initCCDB(TCollision const& collision) + { + int64_t timeStampML = 0; + if constexpr (requires { collision.timestamp(); }) { // we are in derived data + if (mRunNumber == collision.runNumber()) { + return; + } + mRunNumber = collision.runNumber(); + timeStampML = collision.timestamp(); + } + if constexpr (requires { collision.template bc_as(); }) { // we are in original data + auto bc = collision.template bc_as(); + if (mRunNumber == bc.runNumber()) { + return; + } + mRunNumber = bc.runNumber(); + timeStampML = bc.timestamp(); + } + + // machine learning initialization if requested + if (mlConfigurations.calculateXiMinusScores || + mlConfigurations.calculateXiPlusScores || + mlConfigurations.calculateOmegaMinusScores || + mlConfigurations.calculateOmegaPlusScores) { + if (mlConfigurations.timestampCCDB.value != -1) + timeStampML = mlConfigurations.timestampCCDB.value; + LoadMachines(timeStampML); + } + } + + // function to load models for ML-based classifiers + void LoadMachines(int64_t timeStampML) + { + 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?"; + } + } + + 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."; + } + + void init(InitContext const&) + { + // Histograms + histos.add("hEventVertexZ", "hEventVertexZ", kTH1F, {vertexZ}); + + ccdb->setURL(ccdbConfigurations.ccdburl); + } + + // Process candidate and store properties in object + template + void processCandidate(TCascObject const& cand) + { + // Select features + // FIXME THIS NEEDS ADJUSTING + std::vector inputFeatures{0.0f, 0.0f, + 0.0f, 0.0f}; + + // calculate scores + if (mlConfigurations.calculateXiMinusScores && cand.sign() < 0) { + float* xiMinusProbability = mlModelXiMinus.evalModel(inputFeatures); + xiMLSelections(xiMinusProbability[1]); + } + if (mlConfigurations.calculateXiPlusScores && cand.sign() > 0) { + float* xiPlusProbability = mlModelXiPlus.evalModel(inputFeatures); + xiMLSelections(xiPlusProbability[1]); + } + if (mlConfigurations.calculateOmegaMinusScores && cand.sign() < 0) { + float* omegaMinusProbability = mlModelOmegaMinus.evalModel(inputFeatures); + omegaMLSelections(omegaMinusProbability[1]); + } + if (mlConfigurations.calculateOmegaPlusScores && cand.sign() > 0) { + float* omegaPlusProbability = mlModelOmegaPlus.evalModel(inputFeatures); + omegaMLSelections(omegaPlusProbability[1]); + } + } + + void processDerivedData(soa::Join::iterator const& collision, CascDerivedDatas const& cascades) + { + initCCDB(collision); + + histos.fill(HIST("hEventVertexZ"), collision.posZ()); + for (auto& casc : cascades) { + nCandidates++; + if (nCandidates % 50000 == 0) { + LOG(info) << "Candidates processed: " << nCandidates; + } + processCandidate(casc); + } + } + void processStandardData(aod::Collision const& collision, CascOriginalDatas const& cascades) + { + initCCDB(collision); + + histos.fill(HIST("hEventVertexZ"), collision.posZ()); + for (auto& casc : cascades) { + nCandidates++; + if (nCandidates % 50000 == 0) { + LOG(info) << "Candidates processed: " << nCandidates; + } + processCandidate(casc); + } + } + + PROCESS_SWITCH(cascademlselecttion, processStandardData, "Process standard data", false); + PROCESS_SWITCH(cascademlselecttion, processDerivedData, "Process derived data", true); +}; + +WorkflowSpec defineDataProcessing(ConfigContext const& cfgc) +{ + return WorkflowSpec{adaptAnalysisTask(cfgc)}; +} From e7240155532ebbb571b7a1074d3ff782521ce2d7 Mon Sep 17 00:00:00 2001 From: romainschotter Date: Sun, 19 Jan 2025 18:43:16 +0100 Subject: [PATCH 2/4] Add option to store dummy values for ML scores --- .../Strangeness/cascademlselection.cxx | 56 +++++++++++-------- 1 file changed, 34 insertions(+), 22 deletions(-) diff --git a/PWGLF/TableProducer/Strangeness/cascademlselection.cxx b/PWGLF/TableProducer/Strangeness/cascademlselection.cxx index 2add7529ce5..07fb406fdcb 100644 --- a/PWGLF/TableProducer/Strangeness/cascademlselection.cxx +++ b/PWGLF/TableProducer/Strangeness/cascademlselection.cxx @@ -67,7 +67,7 @@ using CascOriginalDatas = soa::Join; // For derived data analysis using CascDerivedDatas = soa::Join; -struct cascademlselecttion { +struct cascademlselection { o2::ml::OnnxModel mlModelXiMinus; o2::ml::OnnxModel mlModelXiPlus; o2::ml::OnnxModel mlModelOmegaMinus; @@ -97,10 +97,10 @@ struct cascademlselecttion { // Machine learning evaluation for pre-selection and corresponding information generation struct : ConfigurableGroup { // ML classifiers: master flags to populate ML Selection tables - Configurable calculateXiMinusScores{"mlConfigurations.calculateXiMinusScores", false, "calculate XiMinus ML scores"}; - Configurable calculateXiPlusScores{"mlConfigurations.calculateXiPlusScores", false, "calculate XiPlus ML scores"}; - Configurable calculateOmegaMinusScores{"mlConfigurations.calculateOmegaMinusScores", false, "calculate OmegaMinus ML scores"}; - Configurable calculateOmegaPlusScores{"mlConfigurations.calculateOmegaPlusScores", false, "calculate OmegaPlus ML scores"}; + 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"}; @@ -231,21 +231,33 @@ struct cascademlselecttion { 0.0f, 0.0f}; // calculate scores - if (mlConfigurations.calculateXiMinusScores && cand.sign() < 0) { - float* xiMinusProbability = mlModelXiMinus.evalModel(inputFeatures); - xiMLSelections(xiMinusProbability[1]); - } - if (mlConfigurations.calculateXiPlusScores && cand.sign() > 0) { - float* xiPlusProbability = mlModelXiPlus.evalModel(inputFeatures); - xiMLSelections(xiPlusProbability[1]); - } - if (mlConfigurations.calculateOmegaMinusScores && cand.sign() < 0) { - float* omegaMinusProbability = mlModelOmegaMinus.evalModel(inputFeatures); - omegaMLSelections(omegaMinusProbability[1]); + 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 (mlConfigurations.calculateOmegaPlusScores && cand.sign() > 0) { - float* omegaPlusProbability = mlModelOmegaPlus.evalModel(inputFeatures); - omegaMLSelections(omegaPlusProbability[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); + } } } @@ -276,11 +288,11 @@ struct cascademlselecttion { } } - PROCESS_SWITCH(cascademlselecttion, processStandardData, "Process standard data", false); - PROCESS_SWITCH(cascademlselecttion, processDerivedData, "Process derived data", true); + PROCESS_SWITCH(cascademlselection, processStandardData, "Process standard data", false); + PROCESS_SWITCH(cascademlselection, processDerivedData, "Process derived data", true); }; WorkflowSpec defineDataProcessing(ConfigContext const& cfgc) { - return WorkflowSpec{adaptAnalysisTask(cfgc)}; + return WorkflowSpec{adaptAnalysisTask(cfgc)}; } From da45e4914937f9732ce40e08f6ba174f3ba8de6e Mon Sep 17 00:00:00 2001 From: romainschotter Date: Sun, 19 Jan 2025 18:49:10 +0100 Subject: [PATCH 3/4] Please consider the following formatting changes --- PWGLF/TableProducer/Strangeness/cascademlselection.cxx | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/PWGLF/TableProducer/Strangeness/cascademlselection.cxx b/PWGLF/TableProducer/Strangeness/cascademlselection.cxx index 07fb406fdcb..6da4bc669e4 100644 --- a/PWGLF/TableProducer/Strangeness/cascademlselection.cxx +++ b/PWGLF/TableProducer/Strangeness/cascademlselection.cxx @@ -75,7 +75,7 @@ struct cascademlselection { std::map metadata; - Produces xiMLSelections; // optionally aggregate information from ML output for posterior analysis (derived data) + 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}; @@ -137,7 +137,7 @@ struct cascademlselection { } mRunNumber = collision.runNumber(); timeStampML = collision.timestamp(); - } + } if constexpr (requires { collision.template bc_as(); }) { // we are in original data auto bc = collision.template bc_as(); if (mRunNumber == bc.runNumber()) { @@ -277,7 +277,7 @@ struct cascademlselection { void processStandardData(aod::Collision const& collision, CascOriginalDatas const& cascades) { initCCDB(collision); - + histos.fill(HIST("hEventVertexZ"), collision.posZ()); for (auto& casc : cascades) { nCandidates++; From 42322fa49aa9c921a2db17f3b4abe71248e6403b Mon Sep 17 00:00:00 2001 From: romainschotter Date: Sun, 19 Jan 2025 18:50:32 +0100 Subject: [PATCH 4/4] Please consider the following formatting changes --- PWGLF/TableProducer/Strangeness/CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/PWGLF/TableProducer/Strangeness/CMakeLists.txt b/PWGLF/TableProducer/Strangeness/CMakeLists.txt index f7759257f73..2aa79c3a555 100644 --- a/PWGLF/TableProducer/Strangeness/CMakeLists.txt +++ b/PWGLF/TableProducer/Strangeness/CMakeLists.txt @@ -141,7 +141,7 @@ o2physics_add_dpl_workflow(lambdakzeromlselection o2physics_add_dpl_workflow(cascademlselection SOURCES cascademlselection.cxx PUBLIC_LINK_LIBRARIES O2Physics::AnalysisCore O2Physics::MLCore - COMPONENT_NAME Analysis) + COMPONENT_NAME Analysis) o2physics_add_dpl_workflow(sigma0builder SOURCES sigma0builder.cxx