ticl::TracksterInferenceByDNN Class Reference

#include <TracksterInferenceByDNN.h>

Inheritance diagram for ticl::TracksterInferenceByDNN:

Public Member Functions
void	inputData (const std::vector< reco::CaloCluster > &layerClusters, std::vector< Trackster > &tracksters) override
void	runInference (std::vector< Trackster > &tracksters) override
	TracksterInferenceByDNN (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_
const std::string	en_modelPath_
const std::string	id_modelPath_
std::vector< std::vector< float > >	input_Data_
std::vector< std::vector< int64_t > >	input_shapes_
const std::vector< std::string >	inputNames_
const cms::Ort::ONNXRuntime *	onnxEnergySession_
const cms::Ort::ONNXRuntime *	onnxPIDSession_
const std::vector< std::string >	output_en_
const std::vector< std::string >	output_id_
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 TracksterInferenceByDNN.h.

Constructor & Destructor Documentation

TracksterInferenceByDNN()

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

explicit

Definition at line 14 of file TracksterInferenceByDNN.cc.

References en_modelPath_, id_modelPath_, onnxEnergySession_, and onnxPIDSession_.

      : TracksterInferenceAlgoBase(conf),
        id_modelPath_(
            conf.getParameter<edm::FileInPath>("onnxPIDModelPath").fullPath()),  // Path to the PID model CLU3D
        en_modelPath_(
            conf.getParameter<edm::FileInPath>("onnxEnergyModelPath").fullPath()),  // Path to the Energy model CLU3D
        inputNames_(conf.getParameter<std::vector<std::string>>("inputNames")),     // Define input names for inference
        output_en_(conf.getParameter<std::vector<std::string>>("output_en")),  // Define output energy for inference
        output_id_(conf.getParameter<std::vector<std::string>>("output_id")),  // Define output PID 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> onnxPIDRuntimeInstance =
        std::make_unique<cms::Ort::ONNXRuntime>(id_modelPath_.c_str());
    onnxPIDSession_ = onnxPIDRuntimeInstance.get();
    static std::unique_ptr<cms::Ort::ONNXRuntime> onnxEnergyRuntimeInstance =
        std::make_unique<cms::Ort::ONNXRuntime>(en_modelPath_.c_str());
    onnxEnergySession_ = onnxEnergyRuntimeInstance.get();
  }

Member Function Documentation

fillPSetDescription()

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

static

Definition at line 128 of file TracksterInferenceByDNN.cc.

References edm::ParameterSetDescription::add().

                                                                                     {
    iDesc.add<int>("algo_verbosity", 0);
    iDesc
        .add<edm::FileInPath>("onnxPIDModelPath",
                              edm::FileInPath("RecoHGCal/TICL/data/ticlv5/onnx_models/patternrecognition/id_v0.onnx"))
        ->setComment("Path to ONNX PID model CLU3D");
    iDesc
        .add<edm::FileInPath>(
            "onnxEnergyModelPath",
            edm::FileInPath("RecoHGCal/TICL/data/ticlv5/onnx_models/patternrecognition/energy_v0.onnx"))
        ->setComment("Path to ONNX Energy model CLU3D");
    iDesc.add<std::vector<std::string>>("inputNames", {"input"});
    iDesc.add<std::vector<std::string>>("output_en", {"enreg_output"});
    iDesc.add<std::vector<std::string>>("output_id", {"pid_output"});
    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", 1);
  }

inputData()

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

overridevirtual

Implements ticl::TracksterInferenceAlgoBase.

Definition at line 39 of file TracksterInferenceByDNN.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::TracksterInferenceByDNN::runInference ( std::vector< Trackster > &  tracksters )

overridevirtual

Implements ticl::TracksterInferenceAlgoBase.

Definition at line 98 of file TracksterInferenceByDNN.cc.

References batchSize_, doPID_, doRegression_, HBHEDarkening_cff::energy, mps_fire::i, input_Data_, input_shapes_, inputNames_, onnxEnergySession_, onnxPIDSession_, output_en_, output_id_, mps_fire::result, cms::Ort::ONNXRuntime::run(), tracksterIndices_, and HitToTracksterAssociation_cfi::tracksters.

                                                                             {
    if (batchSize_ == 0)
      return;  // Exit if no batch

    if (doPID_ and doRegression_) {
      // Run energy model inference
      auto result = onnxEnergySession_->run(inputNames_, input_Data_, input_shapes_, output_en_, batchSize_);
      auto& energyOutputTensor = result[0];
      if (!output_en_.empty()) {
        for (int i = 0; i < static_cast<int>(batchSize_); i++) {
          const float energy = energyOutputTensor[i];
          tracksters[tracksterIndices_[i]].setRegressedEnergy(energy);  // Update energy
        }
      }
    }

    if (doPID_) {
      // Run PID model inference
      auto pidOutput = onnxPIDSession_->run(inputNames_, input_Data_, input_shapes_, output_id_, batchSize_);
      auto pidOutputTensor = pidOutput[0];
      float* probs = pidOutputTensor.data();
      if (!output_id_.empty()) {
        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
        }
      }
    }
  }

Member Data Documentation

batchSize_

int ticl::TracksterInferenceByDNN::batchSize_

private

Definition at line 37 of file TracksterInferenceByDNN.h.

Referenced by inputData(), and runInference().

doPID_

int ticl::TracksterInferenceByDNN::doPID_

private

Definition at line 30 of file TracksterInferenceByDNN.h.

Referenced by runInference().

doRegression_

int ticl::TracksterInferenceByDNN::doRegression_

private

Definition at line 31 of file TracksterInferenceByDNN.h.

Referenced by runInference().

eidMinClusterEnergy_

const float ticl::TracksterInferenceByDNN::eidMinClusterEnergy_

private

Definition at line 26 of file TracksterInferenceByDNN.h.

Referenced by inputData().

eidNClusters_

const int ticl::TracksterInferenceByDNN::eidNClusters_

private

Definition at line 28 of file TracksterInferenceByDNN.h.

Referenced by inputData().

eidNFeatures_

constexpr int ticl::TracksterInferenceByDNN::eidNFeatures_ = 3

staticprivate

Definition at line 29 of file TracksterInferenceByDNN.h.

Referenced by inputData().

eidNLayers_

const int ticl::TracksterInferenceByDNN::eidNLayers_

private

Definition at line 27 of file TracksterInferenceByDNN.h.

Referenced by inputData().

en_modelPath_

const std::string ticl::TracksterInferenceByDNN::en_modelPath_

private

Definition at line 22 of file TracksterInferenceByDNN.h.

Referenced by TracksterInferenceByDNN().

id_modelPath_

const std::string ticl::TracksterInferenceByDNN::id_modelPath_

private

Definition at line 21 of file TracksterInferenceByDNN.h.

Referenced by TracksterInferenceByDNN().

input_Data_

std::vector<std::vector<float> > ticl::TracksterInferenceByDNN::input_Data_

private

Definition at line 36 of file TracksterInferenceByDNN.h.

Referenced by inputData(), and runInference().

input_shapes_

std::vector<std::vector<int64_t> > ticl::TracksterInferenceByDNN::input_shapes_

private

Definition at line 34 of file TracksterInferenceByDNN.h.

Referenced by inputData(), and runInference().

inputNames_

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

private

Definition at line 23 of file TracksterInferenceByDNN.h.

Referenced by runInference().

onnxEnergySession_

const cms::Ort::ONNXRuntime* ticl::TracksterInferenceByDNN::onnxEnergySession_

private

Definition at line 19 of file TracksterInferenceByDNN.h.

Referenced by runInference(), and TracksterInferenceByDNN().

onnxPIDSession_

const cms::Ort::ONNXRuntime* ticl::TracksterInferenceByDNN::onnxPIDSession_

private

Definition at line 18 of file TracksterInferenceByDNN.h.

Referenced by runInference(), and TracksterInferenceByDNN().

output_en_

const std::vector<std::string> ticl::TracksterInferenceByDNN::output_en_

private

Definition at line 24 of file TracksterInferenceByDNN.h.

Referenced by runInference().

output_id_

const std::vector<std::string> ticl::TracksterInferenceByDNN::output_id_

private

Definition at line 25 of file TracksterInferenceByDNN.h.

Referenced by runInference().

rhtools_

hgcal::RecHitTools ticl::TracksterInferenceByDNN::rhtools_

private

Definition at line 33 of file TracksterInferenceByDNN.h.

Referenced by inputData().

tracksterIndices_

std::vector<int> ticl::TracksterInferenceByDNN::tracksterIndices_

private

Definition at line 35 of file TracksterInferenceByDNN.h.

Referenced by inputData(), and runInference().