TICL/plugins/TracksterInferenceByCNNv4.cc

0001 #include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h"
0002 #include "RecoHGCal/TICL/interface/TracksterInferenceByCNNv4.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;  // Use ONNXRuntime namespace
0012
0013   // Constructor for TracksterInferenceByCNNv4
0014   TracksterInferenceByCNNv4::TracksterInferenceByCNNv4(const edm::ParameterSet& conf)
0015       : TracksterInferenceAlgoBase(conf),
0016         modelPath_(conf.getParameter<edm::FileInPath>("onnxModelPath").fullPath()),  // Path to the PID model CLU3D
0017         inputNames_(conf.getParameter<std::vector<std::string>>("inputNames")),      // Define input names for inference
0018         outputNames_(conf.getParameter<std::vector<std::string>>("outputNames")),   // Define output names for inference
0019         eidMinClusterEnergy_(conf.getParameter<double>("eid_min_cluster_energy")),  // Minimum cluster energy
0020         eidNLayers_(conf.getParameter<int>("eid_n_layers")),                        // Number of layers
0021         eidNClusters_(conf.getParameter<int>("eid_n_clusters")),                    // Number of clusters
0022         doPID_(conf.getParameter<int>("doPID")),                                    // Number of clusters
0023         doRegression_(conf.getParameter<int>("doRegression"))                       // Number of clusters
0024   {
0025     // Initialize ONNX Runtime sessions for PID and Energy models
0026     static std::unique_ptr<cms::Ort::ONNXRuntime> onnxRuntimeInstance =
0027         std::make_unique<cms::Ort::ONNXRuntime>(modelPath_.c_str());
0028     onnxSession_ = onnxRuntimeInstance.get();
0029   }
0030
0031   // Method to process input data and prepare it for inference
0032   void TracksterInferenceByCNNv4::inputData(const std::vector<reco::CaloCluster>& layerClusters,
0033                                             std::vector<Trackster>& tracksters) {
0034     tracksterIndices_.clear();  // Clear previous indices
0035     for (int i = 0; i < static_cast<int>(tracksters.size()); i++) {
0036       float sumClusterEnergy = 0.;
0037       for (const unsigned int& vertex : tracksters[i].vertices()) {
0038         sumClusterEnergy += static_cast<float>(layerClusters[vertex].energy());
0039         if (sumClusterEnergy >= eidMinClusterEnergy_) {
0040           tracksters[i].setRegressedEnergy(0.f);  // Set regressed energy to 0
0041           tracksters[i].zeroProbabilities();      // Zero out probabilities
0042           tracksterIndices_.push_back(i);         // Add index to the list
0043           break;
0044         }
0045       }
0046     }
0047
0048     // Prepare input shapes and data for inference
0049     batchSize_ = static_cast<int>(tracksterIndices_.size());
0050     if (batchSize_ == 0)
0051       return;  // Exit if no tracksters
0052
0053     std::vector<int64_t> inputShape = {batchSize_, eidNLayers_, eidNClusters_, eidNFeatures_};
0054     input_shapes_ = {inputShape};
0055
0056     input_Data_.clear();
0057     input_Data_.emplace_back(batchSize_ * eidNLayers_ * eidNClusters_ * eidNFeatures_, 0);
0058
0059     for (int i = 0; i < batchSize_; i++) {
0060       const Trackster& trackster = tracksters[tracksterIndices_[i]];
0061
0062       // Prepare indices and sort clusters based on energy
0063       std::vector<int> clusterIndices(trackster.vertices().size());
0064       for (int k = 0; k < static_cast<int>(trackster.vertices().size()); k++) {
0065         clusterIndices[k] = k;
0066       }
0067
0068       std::sort(clusterIndices.begin(), clusterIndices.end(), [&layerClusters, &trackster](const int& a, const int& b) {
0069         return layerClusters[trackster.vertices(a)].energy() > layerClusters[trackster.vertices(b)].energy();
0070       });
0071
0072       std::vector<int> seenClusters(eidNLayers_, 0);
0073
0074       // Fill input data with cluster information
0075       for (const int& k : clusterIndices) {
0076         const reco::CaloCluster& cluster = layerClusters[trackster.vertices(k)];
0077         int j = rhtools_.getLayerWithOffset(cluster.hitsAndFractions()[0].first) - 1;
0078         if (j < eidNLayers_ && seenClusters[j] < eidNClusters_) {
0079           auto index = (i * eidNLayers_ + j) * eidNFeatures_ * eidNClusters_ + seenClusters[j] * eidNFeatures_;
0080           input_Data_[0][index] =
0081               static_cast<float>(cluster.energy() / static_cast<float>(trackster.vertex_multiplicity(k)));
0082           input_Data_[0][index + 1] = static_cast<float>(std::abs(cluster.eta()));
0083           input_Data_[0][index + 2] = static_cast<float>(cluster.phi());
0084           seenClusters[j]++;
0085         }
0086       }
0087     }
0088   }
0089
0090   // Method to run inference and update tracksters
0091   void TracksterInferenceByCNNv4::runInference(std::vector<Trackster>& tracksters) {
0092     if (batchSize_ == 0)
0093       return;  // Exit if no batch
0094
0095     std::vector<std::vector<float>> outputTensors;
0096     outputTensors = onnxSession_->run(inputNames_, input_Data_, input_shapes_, outputNames_, batchSize_);
0097     if (doPID_ and doRegression_) {
0098       // Run energy model inference
0099       if (!outputNames_.empty()) {
0100         for (int i = 0; i < static_cast<int>(batchSize_); i++) {
0101           const float energy = outputTensors[0][i];
0102           tracksters[tracksterIndices_[i]].setRegressedEnergy(energy);  // Update energy
0103         }
0104       }
0105     }
0106
0107     if (doPID_) {
0108       // Run PID model inference
0109       if (!outputNames_.empty()) {
0110         int probsIdx = outputNames_.empty() ? 0 : 1;
0111         std::vector<float> vec = outputTensors[probsIdx];
0112         float* probs = vec.data();
0113         for (int i = 0; i < batchSize_; i++) {
0114           tracksters[tracksterIndices_[i]].setProbabilities(probs);             // Update probabilities
0115           probs += tracksters[tracksterIndices_[i]].id_probabilities().size();  // Move to next set of probabilities
0116         }
0117       }
0118     }
0119   }
0120
0121   // Method to fill parameter set description for configuration
0122   void TracksterInferenceByCNNv4::fillPSetDescription(edm::ParameterSetDescription& iDesc) {
0123     iDesc.add<int>("algo_verbosity", 0);
0124     iDesc
0125         .add<edm::FileInPath>("onnxModelPath",
0126                               edm::FileInPath("RecoHGCal/TICL/data/ticlv4/onnx_models/energy_id_v0.onnx"))
0127         ->setComment("Path to ONNX PID model CLU3D");
0128     iDesc.add<std::vector<std::string>>("inputNames", {"input:0"});
0129     iDesc.add<std::vector<std::string>>("outputNames", {"output/regressed_energy:0", "output/id_probabilities:0"});
0130     iDesc.add<double>("eid_min_cluster_energy", 1.0);
0131     iDesc.add<int>("eid_n_layers", 50);
0132     iDesc.add<int>("eid_n_clusters", 10);
0133     iDesc.add<int>("doPID", 1);
0134     iDesc.add<int>("doRegression", 0);
0135   }
0136 }  // namespace ticl