#include <TracksterInferenceByCNNv4.h>

Inheritance diagram for ticl::TracksterInferenceByCNNv4:

Public Member Functions
void	inputData (const std::vector< reco::CaloCluster > &layerClusters, std::vector< Trackster > &tracksters) override

void	runInference (std::vector< Trackster > &tracksters) override

	TracksterInferenceByCNNv4 (const edm::ParameterSet &conf)

Public Member Functions inherited from ticl::TracksterInferenceAlgoBase
	TracksterInferenceAlgoBase (const edm::ParameterSet &conf)

virtual	~TracksterInferenceAlgoBase ()

Static Public Member Functions
static void	fillPSetDescription (edm::ParameterSetDescription &iDesc)

Static Public Member Functions inherited from ticl::TracksterInferenceAlgoBase
static void	fillPSetDescription (edm::ParameterSetDescription &desc)

Private Attributes
int	batchSize_

int	doPID_

int	doRegression_

const float	eidMinClusterEnergy_

const int	eidNClusters_

const int	eidNLayers_

std::vector< std::vector< float > >	input_Data_

std::vector< std::vector< int64_t > >	input_shapes_

const std::vector< std::string >	inputNames_

const std::string	modelPath_

const cms::Ort::ONNXRuntime *	onnxSession_

const std::vector< std::string >	outputNames_

hgcal::RecHitTools	rhtools_

std::vector< int >	tracksterIndices_

Static Private Attributes
static constexpr int	eidNFeatures_ = 3

Additional Inherited Members
Protected Attributes inherited from ticl::TracksterInferenceAlgoBase
int	algo_verbosity_

Detailed Description

Definition at line 9 of file TracksterInferenceByCNNv4.h.

Constructor & Destructor Documentation

◆ TracksterInferenceByCNNv4()

ticl::TracksterInferenceByCNNv4::TracksterInferenceByCNNv4 ( const edm::ParameterSet & conf )

explicit

Definition at line 14 of file TracksterInferenceByCNNv4.cc.

References modelPath_, and onnxSession_.

       : TracksterInferenceAlgoBase(conf),
         modelPath_(conf.getParameter<edm::FileInPath>("onnxModelPath").fullPath()),  // Path to the PID model CLU3D
         inputNames_(conf.getParameter<std::vector<std::string>>("inputNames")),      // Define input names for inference
         outputNames_(conf.getParameter<std::vector<std::string>>("outputNames")),   // Define output names for inference
         eidMinClusterEnergy_(conf.getParameter<double>("eid_min_cluster_energy")),  // Minimum cluster energy
         eidNLayers_(conf.getParameter<int>("eid_n_layers")),                        // Number of layers
         eidNClusters_(conf.getParameter<int>("eid_n_clusters")),                    // Number of clusters
         doPID_(conf.getParameter<int>("doPID")),                                    // Number of clusters
         doRegression_(conf.getParameter<int>("doRegression"))                       // Number of clusters
   {
     // Initialize ONNX Runtime sessions for PID and Energy models
     static std::unique_ptr<cms::Ort::ONNXRuntime> onnxRuntimeInstance =
         std::make_unique<cms::Ort::ONNXRuntime>(modelPath_.c_str());
     onnxSession_ = onnxRuntimeInstance.get();
   }

Member Function Documentation

◆ fillPSetDescription()

void ticl::TracksterInferenceByCNNv4::fillPSetDescription ( edm::ParameterSetDescription & iDesc )

static

Definition at line 122 of file TracksterInferenceByCNNv4.cc.

References edm::ParameterSetDescription::add().

                                                                                        {
     iDesc.add<int>("algo_verbosity", 0);
     iDesc
         .add<edm::FileInPath>("onnxModelPath",
                               edm::FileInPath("RecoHGCal/TICL/data/ticlv4/onnx_models/energy_id_v0.onnx"))
         ->setComment("Path to ONNX PID model CLU3D");
     iDesc.add<std::vector<std::string>>("inputNames", {"input:0"});
     iDesc.add<std::vector<std::string>>("outputNames", {"output/regressed_energy:0", "output/id_probabilities:0"});
     iDesc.add<double>("eid_min_cluster_energy", 1.0);
     iDesc.add<int>("eid_n_layers", 50);
     iDesc.add<int>("eid_n_clusters", 10);
     iDesc.add<int>("doPID", 1);
     iDesc.add<int>("doRegression", 0);
   }

◆ inputData()

void ticl::TracksterInferenceByCNNv4::inputData	(	const std::vector< reco::CaloCluster > &	layerClusters,
		std::vector< Trackster > &	tracksters
	)

overridevirtual

Implements ticl::TracksterInferenceAlgoBase.

Definition at line 32 of file TracksterInferenceByCNNv4.cc.

References a, funct::abs(), b, batchSize_, eidMinClusterEnergy_, eidNClusters_, eidNFeatures_, eidNLayers_, reco::CaloCluster::energy(), reco::CaloCluster::eta(), f, hgcal::RecHitTools::getLayerWithOffset(), reco::CaloCluster::hitsAndFractions(), mps_fire::i, input_Data_, input_shapes_, dqmiolumiharvest::j, isotrackApplyRegressor::k, hltEgammaHGCALIDVarsL1Seeded_cfi::layerClusters, reco::CaloCluster::phi(), rhtools_, jetUpdater_cfi::sort, tracksterIndices_, HitToTracksterAssociation_cfi::tracksters, bphysicsOniaDQM_cfi::vertex, ticl::Trackster::vertex_multiplicity(), AlignmentTracksFromVertexSelector_cfi::vertices, and ticl::Trackster::vertices().

                                                                               {
     tracksterIndices_.clear();  // Clear previous indices
     for (int i = 0; i < static_cast<int>(tracksters.size()); i++) {
       float sumClusterEnergy = 0.;
       for (const unsigned int& vertex : tracksters[i].vertices()) {
         sumClusterEnergy += static_cast<float>(layerClusters[vertex].energy());
         if (sumClusterEnergy >= eidMinClusterEnergy_) {
           tracksters[i].setRegressedEnergy(0.f);  // Set regressed energy to 0
           tracksters[i].zeroProbabilities();      // Zero out probabilities
           tracksterIndices_.push_back(i);         // Add index to the list
           break;
         }
       }
     }
 
     // Prepare input shapes and data for inference
     batchSize_ = static_cast<int>(tracksterIndices_.size());
     if (batchSize_ == 0)
       return;  // Exit if no tracksters
 
     std::vector<int64_t> inputShape = {batchSize_, eidNLayers_, eidNClusters_, eidNFeatures_};
     input_shapes_ = {inputShape};
 
     input_Data_.clear();
     input_Data_.emplace_back(batchSize_ * eidNLayers_ * eidNClusters_ * eidNFeatures_, 0);
 
     for (int i = 0; i < batchSize_; i++) {
       const Trackster& trackster = tracksters[tracksterIndices_[i]];
 
       // Prepare indices and sort clusters based on energy
       std::vector<int> clusterIndices(trackster.vertices().size());
       for (int k = 0; k < static_cast<int>(trackster.vertices().size()); k++) {
         clusterIndices[k] = k;
       }
 
       std::sort(clusterIndices.begin(), clusterIndices.end(), [&layerClusters, &trackster](const int& a, const int& b) {
         return layerClusters[trackster.vertices(a)].energy() > layerClusters[trackster.vertices(b)].energy();
       });
 
       std::vector<int> seenClusters(eidNLayers_, 0);
 
       // Fill input data with cluster information
       for (const int& k : clusterIndices) {
         const reco::CaloCluster& cluster = layerClusters[trackster.vertices(k)];
         int j = rhtools_.getLayerWithOffset(cluster.hitsAndFractions()[0].first) - 1;
         if (j < eidNLayers_ && seenClusters[j] < eidNClusters_) {
           auto index = (i * eidNLayers_ + j) * eidNFeatures_ * eidNClusters_ + seenClusters[j] * eidNFeatures_;
           input_Data_[0][index] =
               static_cast<float>(cluster.energy() / static_cast<float>(trackster.vertex_multiplicity(k)));
           input_Data_[0][index + 1] = static_cast<float>(std::abs(cluster.eta()));
           input_Data_[0][index + 2] = static_cast<float>(cluster.phi());
           seenClusters[j]++;
         }
       }
     }
   }

◆ runInference()

void ticl::TracksterInferenceByCNNv4::runInference ( std::vector< Trackster > & tracksters )

overridevirtual

Implements ticl::TracksterInferenceAlgoBase.

Definition at line 91 of file TracksterInferenceByCNNv4.cc.

References batchSize_, doPID_, doRegression_, HBHEDarkening_cff::energy, mps_fire::i, input_Data_, input_shapes_, inputNames_, onnxSession_, outputNames_, cms::Ort::ONNXRuntime::run(), tracksterIndices_, and HitToTracksterAssociation_cfi::tracksters.

                                                                                {
     if (batchSize_ == 0)
       return;  // Exit if no batch
 
     std::vector<std::vector<float>> outputTensors;
     outputTensors = onnxSession_->run(inputNames_, input_Data_, input_shapes_, outputNames_, batchSize_);
     if (doPID_ and doRegression_) {
       // Run energy model inference
       if (!outputNames_.empty()) {
         for (int i = 0; i < static_cast<int>(batchSize_); i++) {
           const float energy = outputTensors[0][i];
           tracksters[tracksterIndices_[i]].setRegressedEnergy(energy);  // Update energy
         }
       }
     }
 
     if (doPID_) {
       // Run PID model inference
       if (!outputNames_.empty()) {
         int probsIdx = outputNames_.empty() ? 0 : 1;
         std::vector<float> vec = outputTensors[probsIdx];
         float* probs = vec.data();
         for (int i = 0; i < batchSize_; i++) {
           tracksters[tracksterIndices_[i]].setProbabilities(probs);             // Update probabilities
           probs += tracksters[tracksterIndices_[i]].id_probabilities().size();  // Move to next set of probabilities
         }
       }
     }
   }