From 70d1cd2adb786c966c9f197194c2d59d087c99f0 Mon Sep 17 00:00:00 2001 From: Christian Sonnabend Date: Mon, 21 Sep 2026 12:45:55 +0200 Subject: [PATCH 1/2] Add CCDB-backed ONNX primary transport pruning --- .../include/SimulationDataFormat/StackParam.h | 8 ++ Detectors/Base/CMakeLists.txt | 1 + Detectors/Base/src/Stack.cxx | 126 +++++++++++++++++- 3 files changed, 133 insertions(+), 2 deletions(-) diff --git a/DataFormats/simulation/include/SimulationDataFormat/StackParam.h b/DataFormats/simulation/include/SimulationDataFormat/StackParam.h index b76112b41b541..75de0f2870ae2 100644 --- a/DataFormats/simulation/include/SimulationDataFormat/StackParam.h +++ b/DataFormats/simulation/include/SimulationDataFormat/StackParam.h @@ -28,6 +28,14 @@ struct StackParam : public o2::conf::ConfigurableParamHelper { std::string transportPrimaryFileName = ""; std::string transportPrimaryFuncName = ""; bool transportPrimaryInvert = false; + // Used when transportPrimary="onnx". The model is fetched as raw ONNX bytes + // and class 1 means "skip GEANT transport". + std::string transportPrimaryOnnxCCDBUrl = "http://alice-ccdb.cern.ch"; + std::string transportPrimaryOnnxCCDBPath = ""; + long transportPrimaryOnnxTimestamp = -1; + float transportPrimaryOnnxThreshold = 0.5f; + int transportPrimaryOnnxOutputIndex = 0; + bool transportPrimaryOnnxApplySigmoid = true; // boilerplate stuff + make principal key "Stack" O2ParamDef(StackParam, "Stack"); diff --git a/Detectors/Base/CMakeLists.txt b/Detectors/Base/CMakeLists.txt index 76e4ed9f741fd..effadfb3b5c12 100644 --- a/Detectors/Base/CMakeLists.txt +++ b/Detectors/Base/CMakeLists.txt @@ -51,6 +51,7 @@ o2_add_library(DetectorsBase O2::SimulationDataFormat O2::SimConfig O2::CCDB + O2::ML O2::GPUDataTypes MC::VMC TBB::tbb diff --git a/Detectors/Base/src/Stack.cxx b/Detectors/Base/src/Stack.cxx index a00c21c0589b9..a9d3398b63424 100644 --- a/Detectors/Base/src/Stack.cxx +++ b/Detectors/Base/src/Stack.cxx @@ -25,23 +25,134 @@ #include "SimulationDataFormat/BaseHits.h" #include "SimulationDataFormat/StackParam.h" #include "CommonUtils/ConfigurationMacroHelper.h" +#include "CCDB/CcdbApi.h" +#include "ML/OrtInterface.h" #include "TLorentzVector.h" // for TLorentzVector #include "TParticle.h" // for TParticle #include "TRefArray.h" // for TRefArray #include "TVirtualMC.h" // for VMC #include "TMCProcess.h" // for VMC Particle Production Process +#include "TParticlePDG.h" #include #include #include // for NULL #include +#include +#include +#include +#include using std::cout; using std::endl; using std::pair; using namespace o2::data; +namespace +{ +// Feature contract used by the sim-pruning models, in order: +// pdg, abs_pdg, charge_sign, mass, energy, ekin, px, py, pz, p, pt, eta, +// phi, theta, rapidity, vx, vy, vz, t_ns, dx/dy/dz_from_event, r_xy, +// r_from_event_xy, r3_from_event. Input normalisation can be embedded in the +// ONNX graph, keeping this code independent of model topology. +constexpr size_t OnnxFeatureCount = 25; + +class OnnxPrimaryTransport +{ + public: + explicit OnnxPrimaryTransport(const o2::sim::StackParam& param) + : mThreshold(param.transportPrimaryOnnxThreshold), + mOutputIndex(param.transportPrimaryOnnxOutputIndex), + mApplySigmoid(param.transportPrimaryOnnxApplySigmoid) + { + if (param.transportPrimaryOnnxCCDBPath.empty()) { + throw std::runtime_error("Stack.transportPrimaryOnnxCCDBPath must be configured"); + } + + o2::ccdb::CcdbApi ccdb; + ccdb.init(param.transportPrimaryOnnxCCDBUrl); + std::map headers; + ccdb.loadFileToMemory(mModelBytes, param.transportPrimaryOnnxCCDBPath, {}, + param.transportPrimaryOnnxTimestamp, &headers, {}, {}, {}); + if (mModelBytes.empty()) { + throw std::runtime_error("failed to retrieve ONNX model from CCDB path " + param.transportPrimaryOnnxCCDBPath); + } + + std::unordered_map options{{"model-path", param.transportPrimaryOnnxCCDBPath}, + {"device-type", "CPU"}, + {"intra-op-num-threads", "1"}, + {"inter-op-num-threads", "1"}, + {"enable-optimizations", "99"}, + {"logging-level", "2"}, + {"onnx-environment-name", "primary-transport-pruning"}}; + mModel.init(options); + mModel.initSessionFromBuffer(mModelBytes.data(), mModelBytes.size()); + + const auto inputShapes = mModel.getNumInputNodes(); + if (inputShapes.size() != 1 || inputShapes[0].empty() || + (inputShapes[0].back() > 0 && inputShapes[0].back() != OnnxFeatureCount)) { + throw std::runtime_error("primary transport ONNX model must have one float input with 25 features"); + } + if (mModel.getNumOutputNodes().size() != 1 || mOutputIndex < 0) { + throw std::runtime_error("primary transport ONNX model must have one output and a non-negative output index"); + } + } + + bool transport(const TParticle& particle, const std::vector& primaries) + { + std::vector> inputs{makeFeatures(particle, primaries)}; + auto output = mModel.inference(inputs); + if (static_cast(mOutputIndex) >= output.size()) { + throw std::runtime_error("Stack.transportPrimaryOnnxOutputIndex is outside the model output"); + } + float score = output[mOutputIndex]; + if (mApplySigmoid) { + score = score >= 0.f ? 1.f / (1.f + std::exp(-score)) : std::exp(score) / (1.f + std::exp(score)); + } + // Class 1 means that GEANT transport can be avoided. + return score < mThreshold; + } + + private: + static std::vector makeFeatures(const TParticle& particle, const std::vector& primaries) + { + const double px = particle.Px(); + const double py = particle.Py(); + const double pz = particle.Pz(); + const double momentum = std::sqrt(px * px + py * py + pz * pz); + const double pt = std::hypot(px, py); + const double mass = particle.GetMass(); + const double energy = std::sqrt(std::max(0., mass * mass + momentum * momentum)); + const double eta = momentum > std::abs(pz) ? 0.5 * std::log((momentum + pz) / (momentum - pz)) : 0.; + const double theta = momentum > 0. ? std::acos(pz / momentum) : 0.; + const double rapidity = energy > std::abs(pz) ? 0.5 * std::log((energy + pz) / (energy - pz)) : 0.; + const auto* pdgInfo = particle.GetPDG(); + const double chargeSign = pdgInfo == nullptr || pdgInfo->Charge() == 0. ? 0. : std::copysign(1., pdgInfo->Charge()); + const TParticle& eventReference = primaries.empty() ? particle : primaries.front(); + const double dx = particle.Vx() - eventReference.Vx(); + const double dy = particle.Vy() - eventReference.Vy(); + const double dz = particle.Vz() - eventReference.Vz(); + const double pdg = particle.GetPdgCode(); + + return {static_cast(pdg), static_cast(std::abs(pdg)), static_cast(chargeSign), + static_cast(mass), static_cast(energy), static_cast(energy - mass), + static_cast(px), static_cast(py), static_cast(pz), static_cast(momentum), + static_cast(pt), static_cast(eta), static_cast(particle.Phi()), static_cast(theta), + static_cast(rapidity), static_cast(particle.Vx()), static_cast(particle.Vy()), + static_cast(particle.Vz()), static_cast(particle.T()), static_cast(dx), + static_cast(dy), static_cast(dz), static_cast(std::hypot(particle.Vx(), particle.Vy())), + static_cast(std::hypot(dx, dy)), static_cast(std::sqrt(dx * dx + dy * dy + dz * dz))}; + } + + o2::ml::OrtModel mModel; + std::vector mModelBytes; // ORT may use model bytes directly; retain them for the session lifetime. + float mThreshold; + int mOutputIndex; + bool mApplySigmoid; +}; +} // namespace + // small helper function to append to vector at arbitrary position template void insertInVector(std::vector& v, I index, T e) @@ -100,16 +211,27 @@ Stack::Stack(Int_t size) transportPrimary = o2::conf::GetFromMacro(param.transportPrimaryFileName, param.transportPrimaryFuncName, "o2::data::Stack::TransportFcn", "stack_transport_primary"); - if (!mTransportPrimary) { + if (!transportPrimary) { LOG(fatal) << "Failed to retrieve external \'transportPrimary\' function: problem with configuration "; } LOG(info) << "Successfully retrieve external \'transportPrimary\' frunction: " << param.transportPrimaryFileName; + } else if (param.transportPrimary.compare("onnx") == 0) { + try { + auto classifier = std::make_shared(param); + transportPrimary = [classifier](const TParticle& p, const std::vector& particles) { + return classifier->transport(p, particles); + }; + LOG(info) << "Successfully configured ONNX primary transport pruning from CCDB path " + << param.transportPrimaryOnnxCCDBPath; + } catch (const std::exception& error) { + LOG(fatal) << "Failed to configure ONNX primary transport pruning: " << error.what(); + } } else { LOG(fatal) << "unsupported \'trasportPrimary\' mode: " << param.transportPrimary; } if (param.transportPrimaryInvert) { - mTransportPrimary = [transportPrimary](const TParticle& p, const std::vector& particles) { return !transportPrimary; }; + mTransportPrimary = [transportPrimary](const TParticle& p, const std::vector& particles) { return !transportPrimary(p, particles); }; } else { mTransportPrimary = transportPrimary; } From 7a19c99a178f357d09ec200c551e9450b9729d39 Mon Sep 17 00:00:00 2001 From: Christian Sonnabend Date: Mon, 21 Sep 2026 12:54:29 +0200 Subject: [PATCH 2/2] Reuse simulation CCDB manager for ONNX model --- .../include/SimulationDataFormat/StackParam.h | 2 -- Detectors/Base/src/Stack.cxx | 12 ++++++++---- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/DataFormats/simulation/include/SimulationDataFormat/StackParam.h b/DataFormats/simulation/include/SimulationDataFormat/StackParam.h index 75de0f2870ae2..8631456a7ba1c 100644 --- a/DataFormats/simulation/include/SimulationDataFormat/StackParam.h +++ b/DataFormats/simulation/include/SimulationDataFormat/StackParam.h @@ -30,9 +30,7 @@ struct StackParam : public o2::conf::ConfigurableParamHelper { bool transportPrimaryInvert = false; // Used when transportPrimary="onnx". The model is fetched as raw ONNX bytes // and class 1 means "skip GEANT transport". - std::string transportPrimaryOnnxCCDBUrl = "http://alice-ccdb.cern.ch"; std::string transportPrimaryOnnxCCDBPath = ""; - long transportPrimaryOnnxTimestamp = -1; float transportPrimaryOnnxThreshold = 0.5f; int transportPrimaryOnnxOutputIndex = 0; bool transportPrimaryOnnxApplySigmoid = true; diff --git a/Detectors/Base/src/Stack.cxx b/Detectors/Base/src/Stack.cxx index a9d3398b63424..dee392b553f12 100644 --- a/Detectors/Base/src/Stack.cxx +++ b/Detectors/Base/src/Stack.cxx @@ -25,7 +25,7 @@ #include "SimulationDataFormat/BaseHits.h" #include "SimulationDataFormat/StackParam.h" #include "CommonUtils/ConfigurationMacroHelper.h" -#include "CCDB/CcdbApi.h" +#include "CCDB/BasicCCDBManager.h" #include "ML/OrtInterface.h" #include "TLorentzVector.h" // for TLorentzVector @@ -70,11 +70,15 @@ class OnnxPrimaryTransport throw std::runtime_error("Stack.transportPrimaryOnnxCCDBPath must be configured"); } - o2::ccdb::CcdbApi ccdb; - ccdb.init(param.transportPrimaryOnnxCCDBUrl); + auto& ccdbManager = o2::ccdb::BasicCCDBManager::instance(); + auto& ccdb = ccdbManager.getCCDBAccessor(); std::map headers; + const auto createdNotAfter = ccdbManager.getCreatedNotAfter(); + const auto createdNotBefore = ccdbManager.getCreatedNotBefore(); ccdb.loadFileToMemory(mModelBytes, param.transportPrimaryOnnxCCDBPath, {}, - param.transportPrimaryOnnxTimestamp, &headers, {}, {}, {}); + ccdbManager.getTimestamp(), &headers, {}, + createdNotAfter ? std::to_string(createdNotAfter) : "", + createdNotBefore ? std::to_string(createdNotBefore) : ""); if (mModelBytes.empty()) { throw std::runtime_error("failed to retrieve ONNX model from CCDB path " + param.transportPrimaryOnnxCCDBPath); }