TrigClusterNNProducer.cxx

#include "Trigger/TrigClusterNNProducer.h"
 
#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstdint>
#include <iomanip>
#include <ios>
#include <numeric>
#include <ratio>
 
#include "Framework/Event.h"
#include "Framework/Exception/Exception.h"
#include "Framework/Logger.h"
#include "Trigger/Event/TrigClusterNNScore.h"
 
namespace trigger {
 
namespace {
 
std::vector<double> logSoftmax(const std::vector<float>& logits) {
  const double max_val = *std::max_element(logits.begin(), logits.end());
  double sum_exp = 0.;
  for (const float logit : logits) sum_exp += std::exp(logit - max_val);
  const double log_sum_exp = max_val + std::log(sum_exp);
 
  std::vector<double> result;
  result.reserve(logits.size());
  for (const float logit : logits) result.push_back(logit - log_sum_exp);
  return result;
}
 
}  // namespace
 
TrigClusterNNProducer::TrigClusterNNProducer(const std::string& name,
                                             framework::Process& process)
    : framework::Producer(name, process),
      data_{std::vector<float>(MAX_CLUSTERS * N_FEATURES, 0.f)} {}
 
void TrigClusterNNProducer::configure(framework::config::Parameters& ps) {
  cluster_coll_name_ = ps.get<std::string>("cluster_coll_name");
  ecal_sum_coll_name_ = ps.get<std::string>("ecal_sum_coll_name");
  hcal_sum_coll_name_ = ps.get<std::string>("hcal_sum_coll_name");
  input_pass_ = ps.get<std::string>("input_pass");
  score_coll_name_ = ps.get<std::string>("score_coll_name");
 
  const auto model_path{ps.get<std::string>("model_path")};
  if (model_path.empty()) {
    EXCEPTION_RAISE("NoModel",
                    "model_path is empty. Pass the path to "
                    "ecal_nn_trigger.onnx.");
  }
  rt_ = std::make_unique<ldmx::ort::ONNXRuntime>(model_path);
 
  // the batch dimension is reported as -1 by the wrapper
  const std::vector<int64_t> expected_shape{-1,
                                            static_cast<int64_t>(N_CLASSES)};
  const auto& outputs = rt_->getOutputNames();
  if (std::find(outputs.begin(), outputs.end(), "logits") == outputs.end() or
      rt_->getOutputShape("logits") != expected_shape) {
    EXCEPTION_RAISE("BadModel", "Model '" + model_path +
                                    "' does not have a 'logits' output of "
                                    "shape (batch, " +
                                    std::to_string(N_CLASSES) + ").");
  }
}
 
void TrigClusterNNProducer::produce(framework::Event& event) {
  const auto start = std::chrono::high_resolution_clock::now();
  nevents_++;
 
  // throws an exception if an input is missing or found in more than
  // one pass; an empty input is valid and the event is still scored
  const auto& clusters{
      event.getCollection<TrigCaloCluster>(cluster_coll_name_, input_pass_)};
  const auto& ecal_sums{
      event.getCollection<TrigEnergySum>(ecal_sum_coll_name_, input_pass_)};
  const auto& hcal_sums{
      event.getCollection<TrigEnergySum>(hcal_sum_coll_name_, input_pass_)};
 
  fillInputs(clusters, ecal_sums, hcal_sums, data_[0]);
  const auto inference_start = std::chrono::high_resolution_clock::now();
  const auto logits{rt_->run({"x"}, data_, {"logits"})[0]};
  inference_time_ +=
      std::chrono::duration<double, std::milli>(
          std::chrono::high_resolution_clock::now() - inference_start)
          .count();
 
  // P(bkg) in double: 1 - P(bkg) saturates to 1 in float
  const double p_bkg = std::exp(logSoftmax(logits)[BKG_CLASS]);
 
  ldmx_log(debug) << "n_clusters = " << clusters.size()
                  << ", P(bkg) = " << p_bkg;
 
  TrigClusterNNScore score;
  score.setLogits(logits);
  score.setPBkg(p_bkg);
  event.add(score_coll_name_, score);
 
  processing_time_ += std::chrono::duration<double, std::milli>(
                          std::chrono::high_resolution_clock::now() - start)
                          .count();
}
 
void TrigClusterNNProducer::onProcessEnd() {
  if (nevents_ == 0) return;
  ldmx_log(info) << "AVG Time/Event: total " << std::fixed
                 << std::setprecision(3) << processing_time_ / nevents_
                 << " ms, inference " << inference_time_ / nevents_ << " ms";
}
 
void TrigClusterNNProducer::fillInputs(
    const TrigCaloClusterCollection& clusters,
    const TrigEnergySumCollection& ecal_sums,
    const TrigEnergySumCollection& hcal_sums, std::vector<float>& inputs) {
  inputs.assign(MAX_CLUSTERS * N_FEATURES, 0.f);
 
  // sort cluster indices by descending energy
  std::vector<std::size_t> order(clusters.size());
  std::iota(order.begin(), order.end(), 0);
  std::stable_sort(order.begin(), order.end(),
                   [&clusters](const std::size_t lhs, const std::size_t rhs) {
                     return clusters[lhs].e() > clusters[rhs].e();
                   });
 
  double ecal_sum_e = 0.;
  for (const auto& sum : ecal_sums) ecal_sum_e += sum.energy();
  // the HCal sums carry hardware (ADC) energy
  int hcal_sum_adc = 0;
  for (const auto& sum : hcal_sums) hcal_sum_adc += sum.hwEnergy();
 
  for (std::size_t slot = 0; slot < MAX_CLUSTERS; ++slot) {
    float* feats = &inputs[slot * N_FEATURES];
    if (slot < order.size()) {
      const auto& cluster = clusters[order[slot]];
      feats[FEAT_X] = cluster.x();
      feats[FEAT_Y] = cluster.y();
      feats[FEAT_Z] = cluster.z();
      feats[FEAT_E] = cluster.e();
      feats[FEAT_DEPTH] = cluster.depth();
      feats[FEAT_N_TP] = cluster.nTP();
    }
    // event-level features are filled in every slot
    feats[FEAT_N_CLUSTERS] = clusters.size();
    feats[FEAT_ECAL_SUM_E] = ecal_sum_e;
    feats[FEAT_HCAL_SUM_ADC] = hcal_sum_adc;
  }
}
 
}  // namespace trigger
 
DECLARE_PRODUCER(trigger::TrigClusterNNProducer);