From 8247afdc32df6b910786f340e3bc5144dfc57936 Mon Sep 17 00:00:00 2001 From: Marek Mytkowski <101755133+mytkom@users.noreply.github.com> Date: Wed, 7 Aug 2024 15:39:34 +0200 Subject: [PATCH] PIDML: evaluate FSE + self-attention network (#7162) * PIDML evaluate FSE + self-attention network (#5) * remove detector count setting and reorder network arguments (with NaNs if detector not available) * update README.md * markdownlint changes * MegaLinter fixes (#6) * fix include missing file, the same way it was before * readd pLimits to ONNXinterface and pass it to ONNXmodel * Please consider the following formatting changes (#9) * improve qaPidML according to new approach --------- Co-authored-by: ALICE Builder --- Tools/PIDML/KaonPidTask.cxx | 3 +- Tools/PIDML/README.md | 65 ++++++++--- Tools/PIDML/pidOnnxInterface.h | 30 ++--- Tools/PIDML/pidOnnxModel.h | 118 ++++++++++++++------ Tools/PIDML/qaPidML.cxx | 78 +++++-------- Tools/PIDML/simpleApplyPidOnnxInterface.cxx | 4 +- Tools/PIDML/simpleApplyPidOnnxModel.cxx | 9 +- 7 files changed, 180 insertions(+), 127 deletions(-) diff --git a/Tools/PIDML/KaonPidTask.cxx b/Tools/PIDML/KaonPidTask.cxx index 19c0972be4a..acc29e47bee 100644 --- a/Tools/PIDML/KaonPidTask.cxx +++ b/Tools/PIDML/KaonPidTask.cxx @@ -60,7 +60,6 @@ struct KaonPidTask { Configurable cfgCCDBURL{"ccdb-url", "http://alice-ccdb.cern.ch", "URL of the CCDB repository"}; Configurable cfgPid{"pid", 321, "PID to predict"}; Configurable cfgCertainty{"certainty", 0.5, "Minimum certainty above which the model accepts a particular type of particle"}; - Configurable cfgDetector{"detector", kTPCTOFTRD, "What detectors to use: 0: TPC only, 1: TPC + TOF, 2: TPC + TOF + TRD"}; Configurable cfgTimestamp{"timestamp", 0, "Fixed timestamp"}; Configurable cfgUseCCDB{"useCCDB", false, "Whether to autofetch ML model from CCDB. If false, local file will be used."}; @@ -85,7 +84,7 @@ struct KaonPidTask { if (cfgUseCCDB) { ccdbApi.init(cfgCCDBURL); // Initializes ccdbApi when cfgUseCCDB is set to 'true' } - pidModel = std::make_shared(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, cfgTimestamp.value, cfgPid.value, static_cast(cfgDetector.value), cfgCertainty.value); + pidModel = std::make_shared(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, cfgTimestamp.value, cfgPid.value, cfgCertainty.value); histos.add("hChargePos", ";z;", kTH1F, {{3, -1.5, 1.5}}); histos.add("hChargeNeg", ";z;", kTH1F, {{3, -1.5, 1.5}}); diff --git a/Tools/PIDML/README.md b/Tools/PIDML/README.md index 57ada2bb48e..1768ac9e900 100644 --- a/Tools/PIDML/README.md +++ b/Tools/PIDML/README.md @@ -1,6 +1,8 @@ # PID ML in O2 -Particle identification is essential in most of the analyzes. The PID ML interface will help you to make use of the machine learning models to improve purity and efficiency of particle kinds for your analysis. A single model is tailored to a specific particle kind, e.g., pions with PID 211. For each track, the model returns a float value in [0, 1] which measures the ''certainty'' of the model that this track is of given kind. +Particle identification is essential in most of the analyzes. +The PID ML interface will help you to make use of the machine learning models to improve purity and efficiency of particle kinds for your analysis. +A single model is tailored to a specific particle kind, e.g., pions with PID 211. For each track, the model returns a float value in [0, 1] which measures the ''certainty'' of the model that this track is of given kind. ## PidONNXModel @@ -11,12 +13,16 @@ This class represents a single ML model from an ONNX file. It requires the follo - CCDB Api instance created in an analysis task - timestamp of the input analysis data -- neded to choose appropriate model - PID to be checked -- detector setup: what detectors should be used for identification. It is described by enum PidMLDetector. Currently available setups: TPC, TPC+TOF, TPC+TOF+TRD - minimum certainty for accepting a track to be of given PID +- *p* limits array - specifiying p limits for each detector configuration (TPC, TPC+TOF, TPC+TOF+TRD) -Let's assume your `PidONNXModel` instance is named `pidModel`. Then, inside your analysis task `process()` function, you can iterate over tracks and call: `pidModel.applyModel(track);` to get the certainty of the model. You can also use `pidModel.applyModelBoolean(track);` to receive a true/false answer, whether the track can be accepted based on the minimum certainty provided to the `PidONNXModel` constructor. +Let's assume your `PidONNXModel` instance is named `pidModel`. +Then, inside your analysis task `process()` function, you can iterate over tracks and call: `pidModel.applyModel(track);` to get the certainty of the model. +You can also use `pidModel.applyModelBoolean(track);` to receive a true/false answer, whether the track can be accepted based on the minimum certainty provided to the `PidONNXModel` constructor. -You can check [a simple analysis task example](https://github.com/AliceO2Group/O2Physics/blob/master/Tools/PIDML/simpleApplyPidOnnxModel.cxx). It uses configurable parameters and shows how to calculate the data timestamp. Note that the calculation of the timestamp requires subscribing to `aod::Collisions` and `aod::BCsWithTimestamps`. For Hyperloop tests, you can set `cfgUseFixedTimestamp` to true with `cfgTimestamp` set to the default value. +You can check [a simple analysis task example](https://github.com/AliceO2Group/O2Physics/blob/master/Tools/PIDML/simpleApplyPidOnnxModel.cxx). +It uses configurable parameters and shows how to calculate the data timestamp. Note that the calculation of the timestamp requires subscribing to `aod::Collisions` and `aod::BCsWithTimestamps`. +For Hyperloop tests, you can set `cfgUseFixedTimestamp` to true with `cfgTimestamp` set to the default value. On the other hand, it is possible to use locally stored models, and then the timestamp is not used, so it can be a dummy value. `processTracksOnly` presents how to analyze on local-only PID ML models. @@ -31,10 +37,10 @@ This is a wrapper around PidONNXModel that contains several models. It has the p Then, obligatory parameters for the interface: - a vector of int output PIDs -- a 2-dimensional LabeledArray of *p*T limits for each PID, for each detector configuration. It describes the minimum *p*T values at which each next detector should be included for predicting given PID +- a 2-dimensional LabeledArray of *p* limits for each PID, for each detector configuration. It describes the minimum *p* values at which each next detector should be included for predicting given PID - a vector of minimum certainties for each PID for accepting a track to be of this PID - boolean flag: whether to switch on auto mode. If true, then *p*T limits and minimum certainties can be passed as an empty array and an empty vector, and the interface will fill them with default configuration: - - *p*T limits: same values for all PIDs: 0.0 (TPC), 0.5 (TPC + TOF), 0.8 (TPC + TOF + TRD) + - *p* limits: same values for all PIDs: 0.0 (TPC), 0.5 (TPC + TOF), 0.8 (TPC + TOF + TRD) - minimum certainties: 0.5 for all PIDs You can use the interface in the same way as the model, by calling `applyModel(track)` or `applyModelBoolean(track)`. The interface will then call the respective method of the model selected with the aforementioned interface parameters. @@ -48,20 +54,49 @@ There is again [a simple analysis task example](https://github.com/AliceO2Group/ Currently, only models for run 285064 (timestamp interval: 1524176895000 - 1524212953000) are uploaded to CCDB, so you can use hardcoded timestamp 1524176895000 for tests. Both model and interface analysis examples can be run with a script: + +### Script for Run2 Converted to Run3 data +```bash +#!/bin/bash + +config_file="my-config.json" + +o2-analysis-tracks-extra-converter --configuration json://$config_file -b | + o2-analysis-timestamp --configuration json://$config_file -b | + o2-analysis-trackextension --configuration json://$config_file -b | + o2-analysis-trackselection --configuration json://$config_file -b | + o2-analysis-multiplicity-table --configuration json://$config_file -b | + o2-analysis-bc-converter --configuration json://$config_file -b | + o2-analysis-collision-converter --configuration json://$config_file -b | + o2-analysis-zdc-converter --configuration json://$config_file -b | + o2-analysis-pid-tof-base --configuration json://$config_file -b | + o2-analysis-pid-tof-beta --configuration json://$config_file -b | + o2-analysis-pid-tof-full --configuration json://$config_file -b | + o2-analysis-pid-tpc-full --configuration json://$config_file -b | + o2-analysis-pid-tpc-base --configuration json://$config_file -b | + o2-analysis-simple-apply-pid-onnx-model --configuration json://$config_file -b +``` +Remember to set every setting, which states that helper task should process Run2 data to `true`. + +### Script for Run3 data ```bash #!/bin/bash config_file="my-config.json" o2-analysis-timestamp --configuration json://$config_file -b | - o2-analysis-trackextension --configuration json://$config_file -b | - o2-analysis-trackselection --configuration json://$config_file -b | - o2-analysis-multiplicity-table --configuration json://$config_file -b | - o2-analysis-fdd-converter --configuration json://$config_file -b | - o2-analysis-pid-tof-base --configuration json://$config_file -b | - o2-analysis-pid-tof-beta --configuration json://$config_file -b | - o2-analysis-pid-tof-full --configuration json://$config_file -b | - o2-analysis-pid-tpc-full --configuration json://$config_file -b | - o2-analysis-simple-apply-pid-onnx-model --configuration json://$config_file -b + o2-analysis-event-selection --configuration json://$config_file -b | + o2-analysis-trackselection --configuration json://$config_file -b | + o2-analysis-multiplicity-table --configuration json://$config_file -b | + o2-analysis-track-propagation --configuration json://$config_file -b | + o2-analysis-pid-tof-base --configuration json://$config_file -b | + o2-analysis-pid-tof-beta --configuration json://$config_file -b | + o2-analysis-pid-tof-full --configuration json://$config_file -b | + o2-analysis-pid-tpc-full --configuration json://$config_file -b | + o2-analysis-pid-tpc-base --configuration json://$config_file -b | + o2-analysis-simple-apply-pid-onnx-model --configuration json://$config_file -b ``` +Remember to set every setting, which states that helper task should process Run3 data to `true`. + + Replace "model" with "interface" in the last line if you want to run the interface workflow. diff --git a/Tools/PIDML/pidOnnxInterface.h b/Tools/PIDML/pidOnnxInterface.h index 0cf8b473d4a..03e7bc19146 100644 --- a/Tools/PIDML/pidOnnxInterface.h +++ b/Tools/PIDML/pidOnnxInterface.h @@ -36,17 +36,17 @@ auto certainties_v = std::vector{certainties, certainties + nPids}; // default values for the cuts constexpr double cuts[nPids][nCutVars] = {{0.0, 0.5, 0.8}, {0.0, 0.5, 0.8}, {0.0, 0.5, 0.8}, {0.0, 0.5, 0.8}, {0.0, 0.5, 0.8}, {0.0, 0.5, 0.8}}; - // row labels static const std::vector pidLabels = { "211", "321", "2212", "0211", "0321", "02212"}; // column labels static const std::vector cutVarLabels = { "TPC", "TPC + TOF", "TPC + TOF + TRD"}; + } // namespace pidml_pt_cuts struct PidONNXInterface { - PidONNXInterface(std::string& localPath, std::string& ccdbPath, bool useCCDB, o2::ccdb::CcdbApi& ccdbApi, uint64_t timestamp, std::vector const& pids, o2::framework::LabeledArray const& pTLimits, std::vector const& minCertainties, bool autoMode) : mNPids{pids.size()}, mPTLimits{pTLimits} + PidONNXInterface(std::string& localPath, std::string& ccdbPath, bool useCCDB, o2::ccdb::CcdbApi& ccdbApi, uint64_t timestamp, std::vector const& pids, o2::framework::LabeledArray const& pLimits, std::vector const& minCertainties, bool autoMode) : mNPids{pids.size()}, mPLimits{pLimits} { if (pids.size() == 0) { LOG(fatal) << "PID ML Interface needs at least 1 output pid to predict"; @@ -54,7 +54,7 @@ struct PidONNXInterface { std::set tmp; for (auto& pid : pids) { if (!tmp.insert(pid).second) { - LOG(fatal) << "PID M Interface: output pids cannot repeat!"; + LOG(fatal) << "PID ML Interface: output pids cannot repeat!"; } } @@ -68,9 +68,7 @@ struct PidONNXInterface { minCertaintiesFilled = minCertainties; } for (std::size_t i = 0; i < mNPids; i++) { - for (uint32_t j = 0; j < kNDetectors; j++) { - mModels.emplace_back(localPath, ccdbPath, useCCDB, ccdbApi, timestamp, pids[i], (PidMLDetector)(kTPCOnly + j), minCertaintiesFilled[i]); - } + mModels.emplace_back(localPath, ccdbPath, useCCDB, ccdbApi, timestamp, pids[i], minCertaintiesFilled[i], mPLimits[i]); } } PidONNXInterface() = default; @@ -84,12 +82,8 @@ struct PidONNXInterface { float applyModel(const T& track, int pid) { for (std::size_t i = 0; i < mNPids; i++) { - if (mModels[i * kNDetectors].mPid == pid) { - for (uint32_t j = 0; j < kNDetectors; j++) { - if (track.pt() >= mPTLimits[i][j] && (j == kNDetectors - 1 || track.pt() < mPTLimits[i][j + 1])) { - return mModels[i * kNDetectors + j].applyModel(track); - } - } + if (mModels[i].mPid == pid) { + return mModels[i].applyModel(track); } } LOG(error) << "No suitable PID ML model found for track: " << track.globalIndex() << " from collision: " << track.collision().globalIndex() << " and expected pid: " << pid; @@ -100,12 +94,8 @@ struct PidONNXInterface { bool applyModelBoolean(const T& track, int pid) { for (std::size_t i = 0; i < mNPids; i++) { - if (mModels[i * kNDetectors].mPid == pid) { - for (uint32_t j = 0; j < kNDetectors; j++) { - if (track.pt() >= mPTLimits[i][j] && (j == kNDetectors - 1 || track.pt() < mPTLimits[i][j + 1])) { - return mModels[i * kNDetectors + j].applyModelBoolean(track); - } - } + if (mModels[i].mPid == pid) { + return mModels[i].applyModelBoolean(track); } } LOG(error) << "No suitable PID ML model found for track: " << track.globalIndex() << " from collision: " << track.collision().globalIndex() << " and expected pid: " << pid; @@ -116,12 +106,12 @@ struct PidONNXInterface { void fillDefaultConfiguration(std::vector& minCertainties) { // FIXME: A more sophisticated strategy should be based on pid values as well - mPTLimits = o2::framework::LabeledArray{pidml_pt_cuts::cuts[0], pidml_pt_cuts::nPids, pidml_pt_cuts::nCutVars, pidml_pt_cuts::pidLabels, pidml_pt_cuts::cutVarLabels}; + mPLimits = o2::framework::LabeledArray{pidml_pt_cuts::cuts[0], pidml_pt_cuts::nPids, pidml_pt_cuts::nCutVars, pidml_pt_cuts::pidLabels, pidml_pt_cuts::cutVarLabels}; minCertainties = std::vector(mNPids, 0.5); } std::vector mModels; std::size_t mNPids; - o2::framework::LabeledArray mPTLimits; + o2::framework::LabeledArray mPLimits; }; #endif // TOOLS_PIDML_PIDONNXINTERFACE_H_ diff --git a/Tools/PIDML/pidOnnxModel.h b/Tools/PIDML/pidOnnxModel.h index 8764496c25c..7700801017a 100644 --- a/Tools/PIDML/pidOnnxModel.h +++ b/Tools/PIDML/pidOnnxModel.h @@ -17,6 +17,9 @@ #ifndef TOOLS_PIDML_PIDONNXMODEL_H_ #define TOOLS_PIDML_PIDONNXMODEL_H_ +#include +#include +#include #include #include #include @@ -41,6 +44,13 @@ enum PidMLDetector { kNDetectors ///< number of available detectors configurations }; +namespace pidml_pt_cuts +{ +// TODO: for now first limit wouldn't be used, +// network needs TPC, so we can either do not cut it by p or return 0.0f as prediction +constexpr double defaultModelPLimits[kNDetectors] = {0.0, 0.5, 0.8}; +} // namespace pidml_pt_cuts + // TODO: Copied from cefpTask, shall we put it in some common utils code? namespace { @@ -63,8 +73,10 @@ bool readJsonFile(const std::string& config, rapidjson::Document& d) struct PidONNXModel { public: - PidONNXModel(std::string& localPath, std::string& ccdbPath, bool useCCDB, o2::ccdb::CcdbApi& ccdbApi, uint64_t timestamp, int pid, PidMLDetector detector, double minCertainty) : mDetector(detector), mPid(pid), mMinCertainty(minCertainty) + PidONNXModel(std::string& localPath, std::string& ccdbPath, bool useCCDB, o2::ccdb::CcdbApi& ccdbApi, uint64_t timestamp, int pid, double minCertainty, const double* pLimits = &pidml_pt_cuts::defaultModelPLimits[0]) : mPid(pid), mMinCertainty(minCertainty), mPLimits(pLimits, pLimits + kNDetectors) { + assert(mPLimits.size() == kNDetectors); + std::string modelFile; loadInputFiles(localPath, ccdbPath, useCCDB, ccdbApi, timestamp, pid, modelFile); @@ -131,27 +143,21 @@ struct PidONNXModel { return getModelOutput(track) >= mMinCertainty; } - PidMLDetector mDetector; int mPid; double mMinCertainty; private: void getModelPaths(std::string const& path, std::string& modelDir, std::string& modelFile, std::string& modelPath, int pid, std::string const& ext) { - modelDir = path + "/TPC"; - if (mDetector >= kTPCTOF) { - modelDir += "_TOF"; - } - if (mDetector >= kTPCTOFTRD) { - modelDir += "_TRD"; - } + modelDir = path; + modelFile = "attention_model_"; - modelFile = "simple_model_"; if (pid < 0) { modelFile += "0" + std::to_string(-pid); } else { modelFile += std::to_string(pid); } + modelFile += ext; modelPath = modelDir + "/" + modelFile; } @@ -203,60 +209,99 @@ struct PidONNXModel { } } + bool almostEqual(float a, float b, float eps = 1e-6f) + { + return std::abs(a - b) <= eps; + } + + template + bool trdMissing(const T& track) + { + static constexpr float kTRDMissingSignal = 0.0f; + + return almostEqual(track.trdSignal(), kTRDMissingSignal); + } + + template + bool tofMissing(const T& track) + { + static constexpr float kEpsilon = 1e-4f; + static constexpr float kTOFMissingSignal = -999.0f; + static constexpr float kTOFMissingBeta = -999.0f; + + // Because of run3 data we use also TOF beta value to determine if signal is present + return almostEqual(track.tofSignal(), kTOFMissingSignal, kEpsilon) || almostEqual(track.beta(), kTOFMissingBeta, kEpsilon); + } + + template + bool inPLimit(const T& track, PidMLDetector detector) + { + return track.p() >= mPLimits[detector]; + } + template std::vector createInputsSingle(const T& track) { // TODO: Hardcoded for now. Planning to implement RowView extension to get runtime access to selected columns // sign is short, trackType and tpcNClsShared uint8_t - float scaledX = (track.x() - mScalingParams.at("fX").first) / mScalingParams.at("fX").second; - float scaledY = (track.y() - mScalingParams.at("fY").first) / mScalingParams.at("fY").second; - float scaledZ = (track.z() - mScalingParams.at("fZ").first) / mScalingParams.at("fZ").second; - float scaledAlpha = (track.alpha() - mScalingParams.at("fAlpha").first) / mScalingParams.at("fAlpha").second; - float scaledTPCNClsShared = (static_cast(track.tpcNClsShared()) - mScalingParams.at("fTPCNClsShared").first) / mScalingParams.at("fTPCNClsShared").second; - float scaledDcaXY = (track.dcaXY() - mScalingParams.at("fDcaXY").first) / mScalingParams.at("fDcaXY").second; - float scaledDcaZ = (track.dcaZ() - mScalingParams.at("fDcaZ").first) / mScalingParams.at("fDcaZ").second; float scaledTPCSignal = (track.tpcSignal() - mScalingParams.at("fTPCSignal").first) / mScalingParams.at("fTPCSignal").second; - std::vector inputValues{track.px(), track.py(), track.pz(), static_cast(track.sign()), scaledX, scaledY, scaledZ, scaledAlpha, static_cast(track.trackType()), scaledTPCNClsShared, scaledDcaXY, scaledDcaZ, track.p(), scaledTPCSignal}; + std::vector inputValues{scaledTPCSignal}; - if (mDetector >= kTPCTOF) { + // When TRD Signal shouldn't be used we pass quiet_NaNs to the network + if (!inPLimit(track, kTPCTOFTRD) || trdMissing(track)) { + inputValues.push_back(std::numeric_limits::quiet_NaN()); + inputValues.push_back(std::numeric_limits::quiet_NaN()); + } else { + float scaledTRDSignal = (track.trdSignal() - mScalingParams.at("fTRDSignal").first) / mScalingParams.at("fTRDSignal").second; + inputValues.push_back(scaledTRDSignal); + inputValues.push_back(track.trdPattern()); + } + + // When TOF Signal shouldn't be used we pass quiet_NaNs to the network + if (!inPLimit(track, kTPCTOF) || tofMissing(track)) { + inputValues.push_back(std::numeric_limits::quiet_NaN()); + inputValues.push_back(std::numeric_limits::quiet_NaN()); + } else { float scaledTOFSignal = (track.tofSignal() - mScalingParams.at("fTOFSignal").first) / mScalingParams.at("fTOFSignal").second; float scaledBeta = (track.beta() - mScalingParams.at("fBeta").first) / mScalingParams.at("fBeta").second; inputValues.push_back(scaledTOFSignal); inputValues.push_back(scaledBeta); } - if (mDetector >= kTPCTOFTRD) { - float scaledTRDSignal = (track.trdSignal() - mScalingParams.at("fTRDSignal").first) / mScalingParams.at("fTRDSignal").second; - float scaledTRDPattern = (track.trdPattern() - mScalingParams.at("fTRDPattern").first) / mScalingParams.at("fTRDPattern").second; - inputValues.push_back(scaledTRDSignal); - inputValues.push_back(scaledTRDPattern); - } + float scaledX = (track.x() - mScalingParams.at("fX").first) / mScalingParams.at("fX").second; + float scaledY = (track.y() - mScalingParams.at("fY").first) / mScalingParams.at("fY").second; + float scaledZ = (track.z() - mScalingParams.at("fZ").first) / mScalingParams.at("fZ").second; + float scaledAlpha = (track.alpha() - mScalingParams.at("fAlpha").first) / mScalingParams.at("fAlpha").second; + float scaledTPCNClsShared = (static_cast(track.tpcNClsShared()) - mScalingParams.at("fTPCNClsShared").first) / mScalingParams.at("fTPCNClsShared").second; + float scaledDcaXY = (track.dcaXY() - mScalingParams.at("fDcaXY").first) / mScalingParams.at("fDcaXY").second; + float scaledDcaZ = (track.dcaZ() - mScalingParams.at("fDcaZ").first) / mScalingParams.at("fDcaZ").second; - return inputValues; - } + inputValues.insert(inputValues.end(), {track.p(), track.pt(), track.px(), track.py(), track.pz(), static_cast(track.sign()), scaledX, scaledY, scaledZ, scaledAlpha, static_cast(track.trackType()), scaledTPCNClsShared, scaledDcaXY, scaledDcaZ}); - // FIXME: Temporary solution, new networks will have sigmoid layer added - float sigmoid(float x) - { - float value = std::max(-100.0f, std::min(100.0f, x)); - return 1.0f / (1.0f + std::exp(-value)); + return inputValues; } template float getModelOutput(const T& track) { + // First rank of the expected model input is -1 which means that it is dynamic axis. + // Axis is exported as dynamic to make it possible to run model inference with the batch of + // tracks at once in the future (batch would need to have the same amount of quiet_NaNs in each row). + // For now we hardcode 1. + static constexpr int64_t batch_size = 1; auto input_shape = mInputShapes[0]; + input_shape[0] = batch_size; + std::vector inputTensorValues = createInputsSingle(track); std::vector inputTensors; + #if __has_include() inputTensors.emplace_back(Ort::Experimental::Value::CreateTensor(inputTensorValues.data(), inputTensorValues.size(), input_shape)); #else - Ort::MemoryInfo mem_info = - Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault); + Ort::MemoryInfo mem_info = Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemType::OrtMemTypeDefault); inputTensors.emplace_back(Ort::Value::CreateTensor(mem_info, inputTensorValues.data(), inputTensorValues.size(), input_shape.data(), input_shape.size())); - #endif // Double-check the dimensions of the input tensor @@ -285,7 +330,7 @@ struct PidONNXModel { LOG(debug) << "output tensor shape: " << printShape(outputTensors[0].GetTensorTypeAndShapeInfo().GetShape()); const float* output_value = outputTensors[0].GetTensorData(); - float certainty = sigmoid(*output_value); // FIXME: Temporary, sigmoid will be added as network layer + float certainty = *output_value; return certainty; } catch (const Ort::Exception& exception) { LOG(error) << "Error running model inference: " << exception.what(); @@ -314,6 +359,7 @@ struct PidONNXModel { std::shared_ptr mSession = nullptr; #endif + std::vector mPLimits; std::vector mInputNames; std::vector> mInputShapes; std::vector mOutputNames; diff --git a/Tools/PIDML/qaPidML.cxx b/Tools/PIDML/qaPidML.cxx index accd055222c..2efd8212f55 100644 --- a/Tools/PIDML/qaPidML.cxx +++ b/Tools/PIDML/qaPidML.cxx @@ -13,6 +13,8 @@ /// \author Łukasz Sawicki /// \since +#include + #include "Framework/runDataProcessing.h" #include "Framework/AnalysisTask.h" #include "Framework/HistogramRegistry.h" @@ -21,7 +23,6 @@ #include "Common/DataModel/PIDResponse.h" #include #include "Tools/PIDML/pidOnnxModel.h" -#include using namespace o2; using namespace o2::framework; @@ -47,9 +48,9 @@ struct pidml { // available particles: 211, 2212, 321 static constexpr int particlesPdgCode[numParticles] = {211, 2212, 321}; - // values of track momentum when to switch from only TPC signal to combined TPC and TOF signal + // values of track momentum when to switch from only TPC signal to combined TPC and TOF signal and to TPC+TOF+TRD // i-th momentum corresponds to the i-th particle - static constexpr float pSwitchValue[numParticles] = {0.5, 0.8, 0.5}; + static constexpr double pSwitchValue[numParticles][kNDetectors] = {{0.0, 0.5, 0.8}, {0.0, 0.8, 0.8}, {0.0, 0.5, 0.8}}; HistogramRegistry histReg{ "allHistograms", @@ -244,34 +245,35 @@ struct pidml { template void fillMcHistos(const T& track, const int pdgCode) { - // pions if (pdgCode == 211) { + // pions histReg.fill(HIST("MC/211"), track.pt()); } else if (pdgCode == -211) { + // antipions histReg.fill(HIST("MC/0211"), track.pt()); - } - // protons - else if (pdgCode == 2212) { + } else if (pdgCode == 2212) { + // protons histReg.fill(HIST("MC/2212"), track.pt()); } else if (pdgCode == -2212) { + // antiprotons histReg.fill(HIST("MC/02212"), track.pt()); - } - // kaons - else if (pdgCode == 321) { + } else if (pdgCode == 321) { + // kaons histReg.fill(HIST("MC/321"), track.pt()); } else if (pdgCode == -321) { + // antikaons histReg.fill(HIST("MC/0321"), track.pt()); - } - // electrons - else if (pdgCode == 11) { + } else if (pdgCode == 11) { + // electrons histReg.fill(HIST("MC/11"), track.pt()); } else if (pdgCode == -11) { + // positrons histReg.fill(HIST("MC/011"), track.pt()); - } - // muons - else if (pdgCode == 13) { + } else if (pdgCode == 13) { + // muons histReg.fill(HIST("MC/13"), track.pt()); } else if (pdgCode == -13) { + // antimuons histReg.fill(HIST("MC/013"), track.pt()); } else { histReg.fill(HIST("MC/else"), track.pt()); @@ -326,15 +328,9 @@ struct pidml { void pidML(const T& track, const int pdgCodeMC) { float pidCertainties[3]; - if (track.p() < pSwitchValue[i]) { - pidCertainties[0] = model211TPC.applyModel(track); - pidCertainties[1] = model2212TPC.applyModel(track); - pidCertainties[2] = model321TPC.applyModel(track); - } else { - pidCertainties[0] = model211All.applyModel(track); - pidCertainties[1] = model2212All.applyModel(track); - pidCertainties[2] = model321All.applyModel(track); - } + pidCertainties[0] = model211.applyModel(track); + pidCertainties[1] = model2212.applyModel(track); + pidCertainties[2] = model321.applyModel(track); int pid = getParticlePdg(pidCertainties); // condition for sign: we want to work only with pi, p and K, without antiparticles if (pid == particlesPdgCode[i] && track.sign() == 1) { @@ -346,14 +342,10 @@ struct pidml { } } - // one model for one particle; Model with all TPC and TOF signal - PidONNXModel model211All; - PidONNXModel model2212All; - PidONNXModel model321All; - // Model with only TPC signal model - PidONNXModel model211TPC; - PidONNXModel model2212TPC; - PidONNXModel model321TPC; + // one model for one particle + PidONNXModel model211; + PidONNXModel model2212; + PidONNXModel model321; Configurable cfgPathCCDB{"ccdb-path", "Users/m/mkabus/PIDML", "base path to the CCDB directory with ONNX models"}; Configurable cfgCCDBURL{"ccdb-url", "http://alice-ccdb.cern.ch", "URL of the CCDB repository"}; @@ -368,13 +360,9 @@ struct pidml { if (cfgUseCCDB) { ccdbApi.init(cfgCCDBURL); } else { - model211All = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1, 211, kTPCTOF, 0.5f); - model2212All = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1, 2211, kTPCTOF, 0.5f); - model321All = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1, 321, kTPCTOF, 0.5f); - - model211TPC = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1, 211, kTPCOnly, 0.5f); - model2212TPC = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1, 2211, kTPCOnly, 0.5f); - model321TPC = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1, 321, kTPCOnly, 0.5f); + model211 = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1, 211, 0.5f, pSwitchValue[0]); + model2212 = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1, 2211, 0.5f, pSwitchValue[1]); + model321 = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1, 321, 0.5f, pSwitchValue[2]); } } @@ -384,13 +372,9 @@ struct pidml { { auto bc = collisions.iteratorAt(0).bc_as(); if (cfgUseCCDB && bc.runNumber() != currentRunNumber) { - model211All = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, bc.timestamp(), 211, kTPCTOF, 0.5f); - model2212All = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, bc.timestamp(), 2211, kTPCTOF, 0.5f); - model321All = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, bc.timestamp(), 321, kTPCTOF, 0.5f); - - model211TPC = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, bc.timestamp(), 211, kTPCOnly, 0.5f); - model2212TPC = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, bc.timestamp(), 2211, kTPCOnly, 0.5f); - model321TPC = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, bc.timestamp(), 321, kTPCOnly, 0.5f); + model211 = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, bc.timestamp(), 211, 0.5f, pSwitchValue[0]); + model2212 = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, bc.timestamp(), 2211, 0.5f, pSwitchValue[1]); + model321 = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, bc.timestamp(), 321, 0.5f, pSwitchValue[2]); } for (auto& track : tracks) { diff --git a/Tools/PIDML/simpleApplyPidOnnxInterface.cxx b/Tools/PIDML/simpleApplyPidOnnxInterface.cxx index 85847320787..82f644103ed 100644 --- a/Tools/PIDML/simpleApplyPidOnnxInterface.cxx +++ b/Tools/PIDML/simpleApplyPidOnnxInterface.cxx @@ -14,6 +14,8 @@ /// /// \author Maja Kabus +#include + #include "Framework/runDataProcessing.h" #include "Framework/AnalysisTask.h" #include "CCDB/CcdbApi.h" @@ -21,8 +23,6 @@ #include "Common/DataModel/PIDResponse.h" #include "Tools/PIDML/pidOnnxInterface.h" -#include - using namespace o2; using namespace o2::framework; using namespace o2::framework::expressions; diff --git a/Tools/PIDML/simpleApplyPidOnnxModel.cxx b/Tools/PIDML/simpleApplyPidOnnxModel.cxx index 00dc843ca47..e261c1e10bc 100644 --- a/Tools/PIDML/simpleApplyPidOnnxModel.cxx +++ b/Tools/PIDML/simpleApplyPidOnnxModel.cxx @@ -14,6 +14,8 @@ /// /// \author Maja Kabus +#include + #include "Framework/runDataProcessing.h" #include "Framework/AnalysisTask.h" #include "CCDB/CcdbApi.h" @@ -21,8 +23,6 @@ #include "Common/DataModel/PIDResponse.h" #include "Tools/PIDML/pidOnnxModel.h" -#include - using namespace o2; using namespace o2::framework; using namespace o2::framework::expressions; @@ -40,7 +40,6 @@ DECLARE_SOA_TABLE(MlPidResults, "AOD", "MLPIDRESULTS", o2::soa::Index<>, mlpidre struct SimpleApplyOnnxModel { PidONNXModel pidModel; // One instance per model, e.g., one per each pid to predict - Configurable cfgDetector{"detector", kTPCTOFTRD, "What detectors to use: 0: TPC only, 1: TPC + TOF, 2: TPC + TOF + TRD"}; Configurable cfgPid{"pid", 211, "PID to predict"}; Configurable cfgCertainty{"certainty", 0.5, "Min certainty of the model to accept given particle to be of given kind"}; @@ -68,7 +67,7 @@ struct SimpleApplyOnnxModel { if (cfgUseCCDB) { ccdbApi.init(cfgCCDBURL); } else { - pidModel = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1, cfgPid.value, static_cast(cfgDetector.value), cfgCertainty.value); + pidModel = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, -1, cfgPid.value, cfgCertainty.value); } } @@ -77,7 +76,7 @@ struct SimpleApplyOnnxModel { auto bc = collisions.iteratorAt(0).bc_as(); if (cfgUseCCDB && bc.runNumber() != currentRunNumber) { uint64_t timestamp = cfgUseFixedTimestamp ? cfgTimestamp.value : bc.timestamp(); - pidModel = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, timestamp, cfgPid.value, static_cast(cfgDetector.value), cfgCertainty.value); + pidModel = PidONNXModel(cfgPathLocal.value, cfgPathCCDB.value, cfgUseCCDB.value, ccdbApi, timestamp, cfgPid.value, cfgCertainty.value); } for (auto& track : tracks) {