TICL/plugins/TracksterInferenceByPFN.cc

0001 #include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h"
0002 #include "RecoHGCal/TICL/interface/TracksterInferenceByPFN.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 TracksterInferenceByPFN
0014   TracksterInferenceByPFN::TracksterInferenceByPFN(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")),  // Define input names for inference
0021         output_en_(conf.getParameter<std::vector<std::string>>("output_en")),    // Define output energy for inference
0022         output_id_(conf.getParameter<std::vector<std::string>>("output_id")),    // Define output PID for inference
0023         eidMinClusterEnergy_(conf.getParameter<double>("eid_min_cluster_energy")),  // Minimum cluster energy
0024         eidNLayers_(conf.getParameter<int>("eid_n_layers")),                        // Number of layers
0025         eidNClusters_(conf.getParameter<int>("eid_n_clusters")),                    // Number of clusters
0026         doPID_(conf.getParameter<int>("doPID")),                                    // Number of clusters
0027         doRegression_(conf.getParameter<int>("doRegression"))                       // Number of clusters
0028   {
0029     onnxPIDSession_ = onnxPIDRuntimeInstance_.get();
0030     onnxEnergySession_ = onnxEnergyRuntimeInstance_.get();
0031   }
0032
0033   // Method to process input data and prepare it for inference
0034   void TracksterInferenceByPFN::inputData(const std::vector<reco::CaloCluster>& layerClusters,
0035                                           std::vector<Trackster>& tracksters) {
0036     tracksterIndices_.clear();  // Clear previous indices
0037     for (int i = 0; i < static_cast<int>(tracksters.size()); i++) {
0038       float sumClusterEnergy = 0.;
0039       for (const unsigned int& vertex : tracksters[i].vertices()) {
0040         sumClusterEnergy += static_cast<float>(layerClusters[vertex].energy());
0041         if (sumClusterEnergy >= eidMinClusterEnergy_) {
0042           tracksters[i].setRegressedEnergy(0.f);  // Set regressed energy to 0
0043           tracksters[i].zeroProbabilities();      // Zero out probabilities
0044           tracksterIndices_.push_back(i);         // Add index to the list
0045           break;
0046         }
0047       }
0048     }
0049
0050     // Prepare input shapes and data for inference
0051     batchSize_ = static_cast<int>(tracksterIndices_.size());
0052     if (batchSize_ == 0)
0053       return;  // Exit if no tracksters
0054
0055     std::vector<int64_t> inputShape_lc = {batchSize_, eidNLayers_, eidNClusters_, eidNFeatures_};
0056     std::vector<int64_t> inputShape_tr = {batchSize_, eidNFeatures_};
0057     input_shapes_ = {inputShape_lc, inputShape_tr};
0058
0059     input_Data_.clear();
0060     input_Data_.emplace_back(batchSize_ * eidNLayers_ * eidNClusters_ * eidNFeatures_, 0);
0061     input_Data_.emplace_back(batchSize_ * eidNFeatures_, 0);
0062
0063     for (int i = 0; i < batchSize_; i++) {
0064       const Trackster& trackster = tracksters[tracksterIndices_[i]];
0065       auto index_tr = i * eidNFeatures_;
0066       input_Data_[1][index_tr] = static_cast<float>(trackster.raw_energy());
0067       input_Data_[1][index_tr + 1] = static_cast<float>(trackster.raw_em_energy());
0068       input_Data_[1][index_tr + 2] = static_cast<float>(trackster.barycenter().x());
0069       input_Data_[1][index_tr + 3] = static_cast<float>(trackster.barycenter().y());
0070       input_Data_[1][index_tr + 4] = static_cast<float>(std::abs(trackster.barycenter().z()));
0071       input_Data_[1][index_tr + 5] = static_cast<float>(std::abs(trackster.barycenter().eta()));
0072       input_Data_[1][index_tr + 6] = static_cast<float>(trackster.barycenter().phi());
0073
0074       // Prepare indices and sort clusters based on energy
0075       std::vector<int> clusterIndices(trackster.vertices().size());
0076       for (int k = 0; k < static_cast<int>(trackster.vertices().size()); k++) {
0077         clusterIndices[k] = k;
0078       }
0079
0080       std::sort(clusterIndices.begin(), clusterIndices.end(), [&layerClusters, &trackster](const int& a, const int& b) {
0081         return layerClusters[trackster.vertices(a)].energy() > layerClusters[trackster.vertices(b)].energy();
0082       });
0083
0084       std::vector<int> seenClusters(eidNLayers_, 0);
0085
0086       // Fill input data with cluster information
0087       for (const int& k : clusterIndices) {
0088         const reco::CaloCluster& cluster = layerClusters[trackster.vertices(k)];
0089         int j = rhtools_.getLayerWithOffset(cluster.hitsAndFractions()[0].first) - 1;
0090         if (j < eidNLayers_ && seenClusters[j] < eidNClusters_) {
0091           auto index_lc = (i * eidNLayers_ + j) * eidNFeatures_ * eidNClusters_ + seenClusters[j] * eidNFeatures_;
0092           // Adding more features regarding LC, such as E, eta, phi, x, y, z, and nhits.
0093           input_Data_[0][index_lc] =
0094               static_cast<float>(cluster.energy() / static_cast<float>(trackster.vertex_multiplicity(k)));
0095           input_Data_[0][index_lc + 1] = static_cast<float>(std::abs(cluster.eta()));
0096           input_Data_[0][index_lc + 2] = static_cast<float>(cluster.phi());
0097           input_Data_[0][index_lc + 3] = static_cast<float>(cluster.x());
0098           input_Data_[0][index_lc + 4] = static_cast<float>(cluster.y());
0099           input_Data_[0][index_lc + 5] = static_cast<float>(std::abs(cluster.z()));
0100           input_Data_[0][index_lc + 6] = static_cast<float>(cluster.hitsAndFractions().size());
0101           seenClusters[j]++;
0102         }
0103       }
0104     }
0105   }
0106
0107   // Method to run inference and update tracksters
0108   void TracksterInferenceByPFN::runInference(std::vector<Trackster>& tracksters) {
0109     if (batchSize_ == 0)
0110       return;  // Exit if no batch
0111
0112     if (doPID_ and doRegression_) {
0113       // Run energy model inference
0114       auto result = onnxEnergySession_->run(inputNames_, input_Data_, input_shapes_, output_en_, batchSize_);
0115       auto& energyOutputTensor = result[0];
0116       if (!output_en_.empty()) {
0117         for (int i = 0; i < static_cast<int>(batchSize_); i++) {
0118           const float energy = energyOutputTensor[i];
0119           tracksters[tracksterIndices_[i]].setRegressedEnergy(energy);  // Update energy
0120         }
0121       }
0122     }
0123     if (doPID_) {
0124       // Run PID model inference
0125       auto pidOutput = onnxPIDSession_->run(inputNames_, input_Data_, input_shapes_, output_id_, batchSize_);
0126       auto pidOutputTensor = pidOutput[0];
0127       float* probs = pidOutputTensor.data();
0128       if (!output_id_.empty()) {
0129         for (int i = 0; i < batchSize_; i++) {
0130           tracksters[tracksterIndices_[i]].setProbabilities(probs);             // Update probabilities
0131           probs += tracksters[tracksterIndices_[i]].id_probabilities().size();  // Move to next set of probabilities
0132         }
0133       }
0134     }
0135   }
0136   // Method to fill parameter set description for configuration
0137   void TracksterInferenceByPFN::fillPSetDescription(edm::ParameterSetDescription& iDesc) {
0138     iDesc.add<int>("algo_verbosity", 0);
0139     iDesc
0140         .add<edm::FileInPath>(
0141             "onnxPIDModelPath",
0142             edm::FileInPath("RecoHGCal/TICL/data/ticlv5/onnx_models/PFN/patternrecognition/id_v0.onnx"))
0143         ->setComment("Path to ONNX PID model CLU3D");
0144     iDesc
0145         .add<edm::FileInPath>(
0146             "onnxEnergyModelPath",
0147             edm::FileInPath("RecoHGCal/TICL/data/ticlv5/onnx_models/PFN/patternrecognition/energy_v0.onnx"))
0148         ->setComment("Path to ONNX Energy model CLU3D");
0149     iDesc.add<std::vector<std::string>>("inputNames", {"input", "input_tr_features"});
0150     iDesc.add<std::vector<std::string>>("output_en", {"enreg_output"});
0151     iDesc.add<std::vector<std::string>>("output_id", {"pid_output"});
0152     iDesc.add<double>("eid_min_cluster_energy", 1.0);
0153     iDesc.add<int>("eid_n_layers", 50);
0154     iDesc.add<int>("eid_n_clusters", 10);
0155     iDesc.add<int>("doPID", 1);
0156     iDesc.add<int>("doRegression", 1);
0157   }
0158 }  // namespace ticl