File indexing completed on 2025-07-09 05:00:25
0001 #include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h"
0002 #include "RecoHGCal/TICL/interface/TracksterInferenceByDNN.h"
0003 #include "RecoHGCal/TICL/interface/TracksterInferenceAlgoFactory.h"
0004 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0005 #include "FWCore/Framework/interface/MakerMacros.h"
0006 #include "RecoHGCal/TICL/interface/PatternRecognitionAlgoBase.h"
0007 #include "RecoLocalCalo/HGCalRecAlgos/interface/RecHitTools.h"
0008 #include "TrackstersPCA.h"
0009
0010 namespace ticl {
0011 using namespace cms::Ort;
0012
0013
0014 TracksterInferenceByDNN::TracksterInferenceByDNN(const edm::ParameterSet& conf)
0015 : TracksterInferenceAlgoBase(conf),
0016 onnxPIDRuntimeInstance_(std::make_unique<cms::Ort::ONNXRuntime>(
0017 conf.getParameter<edm::FileInPath>("onnxPIDModelPath").fullPath().c_str())),
0018 onnxEnergyRuntimeInstance_(std::make_unique<cms::Ort::ONNXRuntime>(
0019 conf.getParameter<edm::FileInPath>("onnxEnergyModelPath").fullPath().c_str())),
0020 inputNames_(conf.getParameter<std::vector<std::string>>("inputNames")),
0021 output_en_(conf.getParameter<std::vector<std::string>>("output_en")),
0022 output_id_(conf.getParameter<std::vector<std::string>>("output_id")),
0023 eidMinClusterEnergy_(conf.getParameter<double>("eid_min_cluster_energy")),
0024 eidNLayers_(conf.getParameter<int>("eid_n_layers")),
0025 eidNClusters_(conf.getParameter<int>("eid_n_clusters")),
0026 doPID_(conf.getParameter<int>("doPID")),
0027 doRegression_(conf.getParameter<int>("doRegression"))
0028 {
0029
0030 onnxPIDSession_ = onnxPIDRuntimeInstance_.get();
0031 onnxEnergySession_ = onnxEnergyRuntimeInstance_.get();
0032 }
0033
0034
0035 void TracksterInferenceByDNN::inputData(const std::vector<reco::CaloCluster>& layerClusters,
0036 std::vector<Trackster>& tracksters) {
0037 tracksterIndices_.clear();
0038 for (int i = 0; i < static_cast<int>(tracksters.size()); i++) {
0039 float sumClusterEnergy = 0.;
0040 for (const unsigned int& vertex : tracksters[i].vertices()) {
0041 if (rhtools_.isBarrel(layerClusters[vertex].seed()))
0042 continue;
0043 sumClusterEnergy += static_cast<float>(layerClusters[vertex].energy());
0044 if (sumClusterEnergy >= eidMinClusterEnergy_) {
0045 tracksters[i].setRegressedEnergy(0.f);
0046 tracksters[i].zeroProbabilities();
0047 tracksterIndices_.push_back(i);
0048 break;
0049 }
0050 }
0051 }
0052
0053
0054 batchSize_ = static_cast<int>(tracksterIndices_.size());
0055 if (batchSize_ == 0)
0056 return;
0057
0058 std::vector<int64_t> inputShape = {batchSize_, eidNLayers_, eidNClusters_, eidNFeatures_};
0059 input_shapes_ = {inputShape};
0060
0061 input_Data_.clear();
0062 input_Data_.emplace_back(batchSize_ * eidNLayers_ * eidNClusters_ * eidNFeatures_, 0);
0063
0064 for (int i = 0; i < batchSize_; i++) {
0065 const Trackster& trackster = tracksters[tracksterIndices_[i]];
0066
0067
0068 std::vector<int> clusterIndices(trackster.vertices().size());
0069 for (int k = 0; k < static_cast<int>(trackster.vertices().size()); k++) {
0070 clusterIndices[k] = k;
0071 }
0072
0073 std::sort(clusterIndices.begin(), clusterIndices.end(), [&layerClusters, &trackster](const int& a, const int& b) {
0074 return layerClusters[trackster.vertices(a)].energy() > layerClusters[trackster.vertices(b)].energy();
0075 });
0076
0077 std::vector<int> seenClusters(eidNLayers_, 0);
0078
0079
0080 for (const int& k : clusterIndices) {
0081 const reco::CaloCluster& cluster = layerClusters[trackster.vertices(k)];
0082 int j = rhtools_.getLayerWithOffset(cluster.hitsAndFractions()[0].first) - 1;
0083 if (j < eidNLayers_ && seenClusters[j] < eidNClusters_) {
0084 auto index = (i * eidNLayers_ + j) * eidNFeatures_ * eidNClusters_ + seenClusters[j] * eidNFeatures_;
0085 input_Data_[0][index] =
0086 static_cast<float>(cluster.energy() / static_cast<float>(trackster.vertex_multiplicity(k)));
0087 input_Data_[0][index + 1] = static_cast<float>(std::abs(cluster.eta()));
0088 input_Data_[0][index + 2] = static_cast<float>(cluster.phi());
0089 seenClusters[j]++;
0090 }
0091 }
0092 }
0093 }
0094
0095
0096 void TracksterInferenceByDNN::runInference(std::vector<Trackster>& tracksters) {
0097 if (batchSize_ == 0)
0098 return;
0099
0100 if (doPID_ and doRegression_) {
0101
0102 auto result = onnxEnergySession_->run(inputNames_, input_Data_, input_shapes_, output_en_, batchSize_);
0103 auto& energyOutputTensor = result[0];
0104 if (!output_en_.empty()) {
0105 for (int i = 0; i < static_cast<int>(batchSize_); i++) {
0106 const float energy = energyOutputTensor[i];
0107 tracksters[tracksterIndices_[i]].setRegressedEnergy(energy);
0108 }
0109 }
0110 }
0111
0112 if (doPID_) {
0113
0114 auto pidOutput = onnxPIDSession_->run(inputNames_, input_Data_, input_shapes_, output_id_, batchSize_);
0115 auto pidOutputTensor = pidOutput[0];
0116 float* probs = pidOutputTensor.data();
0117 if (!output_id_.empty()) {
0118 for (int i = 0; i < batchSize_; i++) {
0119 tracksters[tracksterIndices_[i]].setProbabilities(probs);
0120 probs += tracksters[tracksterIndices_[i]].id_probabilities().size();
0121 }
0122 }
0123 }
0124 }
0125
0126 void TracksterInferenceByDNN::fillPSetDescription(edm::ParameterSetDescription& iDesc) {
0127 iDesc.add<int>("algo_verbosity", 0);
0128 iDesc
0129 .add<edm::FileInPath>(
0130 "onnxPIDModelPath",
0131 edm::FileInPath("RecoHGCal/TICL/data/ticlv5/onnx_models/DNN/patternrecognition/id_v0.onnx"))
0132 ->setComment("Path to ONNX PID model CLU3D");
0133 iDesc
0134 .add<edm::FileInPath>(
0135 "onnxEnergyModelPath",
0136 edm::FileInPath("RecoHGCal/TICL/data/ticlv5/onnx_models/DNN/patternrecognition/energy_v0.onnx"))
0137 ->setComment("Path to ONNX Energy model CLU3D");
0138 iDesc.add<std::vector<std::string>>("inputNames", {"input"});
0139 iDesc.add<std::vector<std::string>>("output_en", {"enreg_output"});
0140 iDesc.add<std::vector<std::string>>("output_id", {"pid_output"});
0141 iDesc.add<double>("eid_min_cluster_energy", 1.0);
0142 iDesc.add<int>("eid_n_layers", 50);
0143 iDesc.add<int>("eid_n_clusters", 10);
0144 iDesc.add<int>("doPID", 1);
0145 iDesc.add<int>("doRegression", 1);
0146 }
0147 }