ools/Root/DiTauOnnxDiscriminantTool.cxx

Go to the documentation of this file.

/*
  Copyright (C) 2002-2025 CERN for the benefit of the ATLAS collaboration
*/
 
#include "DiTauRecTools/DiTauOnnxDiscriminantTool.h"
 
// Core include(s):
#include "AthLinks/ElementLink.h"
#include "PathResolver/PathResolver.h"
#include <AsgDataHandles/ReadDecorHandle.h>
 
using TrackParticleLinks_t = std::vector<ElementLink<xAOD::TrackParticleContainer>>;
 
using namespace DiTauRecTools;
 
//=================================PUBLIC-PART==================================
//______________________________________________________________________________
DiTauOnnxDiscriminantTool::DiTauOnnxDiscriminantTool(const std::string& name) 
  : AsgTool(name)
{
}
 
//______________________________________________________________________________
DiTauOnnxDiscriminantTool::~DiTauOnnxDiscriminantTool() = default;
 
//______________________________________________________________________________
StatusCode DiTauOnnxDiscriminantTool::initialize()
{
  ATH_MSG_INFO( "Initializing DiTauOnnxDiscriminantTool" );
  ATH_MSG_INFO( "onnxModelPath: " << m_onnxModelPath );
  std::string model_path = PathResolverFindCalibFile(m_onnxModelPath);
  if (model_path.empty()) {
    ATH_MSG_ERROR("Could not find model file: " << m_onnxModelPath);
    return StatusCode::FAILURE;
  }
  m_ort_env = std::make_unique<Ort::Env>(ORT_LOGGING_LEVEL_WARNING, "OnnxUtil");
  Ort::SessionOptions session_options;
  session_options.SetIntraOpNumThreads(1);
  session_options.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_EXTENDED);
  m_ort_session = std::make_unique<Ort::Session>(*m_ort_env, model_path.c_str(), session_options);
 
  ATH_CHECK( m_ditauContainerKey.initialize() );
  ATH_CHECK( m_ditau_pt_DecorKey.initialize() );
  ATH_CHECK( m_f_core_lead_DecorKey.initialize() );
  ATH_CHECK( m_f_core_sublead_DecorKey.initialize() );
  ATH_CHECK( m_f_subjet_subl_DecorKey.initialize() );
  ATH_CHECK( m_f_subjets_DecorKey.initialize() );
  ATH_CHECK( m_R_max_lead_DecorKey.initialize() );
  ATH_CHECK( m_R_max_sublead_DecorKey.initialize() );
  ATH_CHECK( m_n_track_DecorKey.initialize() );
  ATH_CHECK( m_R_track_all_DecorKey.initialize() );  
  ATH_CHECK( m_R_isotrack_DecorKey.initialize() );
  ATH_CHECK( m_R_track_sublead_DecorKey.initialize() );
  ATH_CHECK( m_M_core_lead_DecorKey.initialize() );
  ATH_CHECK( m_M_core_sublead_DecorKey.initialize() );
  ATH_CHECK( m_M_track_lead_DecorKey.initialize() );
  ATH_CHECK( m_d0_leadtrack_lead_DecorKey.initialize() );
  ATH_CHECK( m_d0_leadtrack_sublead_DecorKey.initialize() );
  ATH_CHECK( m_f_isotracks_DecorKey.initialize() ); 
 
  return StatusCode::SUCCESS;
}
 
StatusCode DiTauOnnxDiscriminantTool::execute(const xAOD::DiTauJet& xDiTau) const
{
    const static SG::Decorator<float> omni_scoreDec("omni_score");
    ATH_MSG_DEBUG("Inferencing omni DiTau ID score...");
    float score = GetDiTauObjOnnxScore(xDiTau);
    ATH_MSG_DEBUG("DiTau ID score: " << score);
    omni_scoreDec(xDiTau) = score;
    return StatusCode::SUCCESS;
}
 
float DiTauOnnxDiscriminantTool::nan_to_num(float value, float nan_replacement = 0.0f, float posinf_replacement = 0.0f, float neginf_replacement = 0.0f) const{
  if (std::isnan(value))
    return nan_replacement;
  if (value == std::numeric_limits<float>::infinity())
    return posinf_replacement;
  if (value == -std::numeric_limits<float>::infinity())
    return neginf_replacement;
  return value;
  }
 
std::vector<float> DiTauOnnxDiscriminantTool::flatten(const std::vector<std::vector<float>> &vec_2d) const{
  std::vector<float> flattened;
  flattened.reserve(vec_2d.size() * (vec_2d.empty() ? 0 : vec_2d[0].size()));
  for (const auto &inner : vec_2d) {
    flattened.insert(flattened.end(), inner.begin(), inner.end());
  }
  return flattened;
}
 
std::vector<float> DiTauOnnxDiscriminantTool::extract_points(const std::vector<std::vector<float>> &track_features) const{
  std::vector<float> points;
  points.reserve(track_features.size() * 2);
  for (const auto &track : track_features) {
    points.push_back(track[0]);  // delta_eta
    points.push_back(track[1]);  // delta_phi
  }
  return points;
}
 
std::vector<float> DiTauOnnxDiscriminantTool::create_mask(const std::vector<std::vector<float>> &track_features) const{
  std::vector<float> mask;
  mask.reserve(track_features.size());
  std::transform(track_features.begin(), track_features.end(), std::back_inserter(mask), [](const auto &track) {
        return std::abs(track[2]) > 1e-6 ? 1.0f : 0.0f;
  });
  return mask;
}
 
Ort::Value DiTauOnnxDiscriminantTool::create_tensor(std::vector<float> &data, const std::vector<int64_t> &shape) const{
  Ort::MemoryInfo memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault);
  return Ort::Value::CreateTensor<float>(memory_info, data.data(), data.size(),shape.data(), shape.size());
}
 
DiTauOnnxDiscriminantTool::InferenceOutput DiTauOnnxDiscriminantTool::run_inference(OnnxInputs &inputs) const{
  std::vector<Ort::Value> input_tensors;
  input_tensors.reserve(m_input_node_names.size());
  input_tensors.emplace_back(create_tensor(inputs.input_features, inputs.input_features_shape));
  input_tensors.emplace_back(create_tensor(inputs.input_points, inputs.input_points_shape));
  input_tensors.emplace_back(create_tensor(inputs.input_mask, inputs.input_mask_shape));
  input_tensors.emplace_back(create_tensor(inputs.input_jet, inputs.input_jet_shape));
  input_tensors.emplace_back(create_tensor(inputs.input_time, inputs.input_time_shape));
 
  std::vector<const char *> input_node_names;
  input_node_names.reserve(m_input_node_names.size());
  std::transform(m_input_node_names.begin(), m_input_node_names.end(), std::back_inserter(input_node_names), [](const std::string &name) { return name.c_str(); });
 
  std::vector<const char *> output_node_names;
  output_node_names.reserve(m_output_node_names.size());
  std::transform(m_output_node_names.begin(), m_output_node_names.end(), std::back_inserter(output_node_names), [](const std::string &name) { return name.c_str(); });
 
  auto output_tensors = m_ort_session->Run(Ort::RunOptions{nullptr}, input_node_names.data(), input_tensors.data(), input_node_names.size(), output_node_names.data(), output_node_names.size());
 
  InferenceOutput output;
  for (size_t i = 0; i < output_tensors.size(); ++i) {
    const auto &tensor = output_tensors[i];
    const size_t length = tensor.GetTensorTypeAndShapeInfo().GetElementCount();
    const float *data = tensor.GetTensorData<float>();
    (i == 0 ? output.output_1 : output.output_2) = std::vector<float>(data, data + length);
  } 
  return output;
}
 
float DiTauOnnxDiscriminantTool::GetDiTauObjOnnxScore(const xAOD::DiTauJet& ditau) const{
 
    SG::ReadDecorHandle<xAOD::DiTauJetContainer,float> ditau_ptDec(m_ditau_pt_DecorKey);    
    SG::ReadDecorHandle<xAOD::DiTauJetContainer,float> f_core_leadDec(m_f_core_lead_DecorKey);
    SG::ReadDecorHandle<xAOD::DiTauJetContainer,float> f_core_sublDec(m_f_core_sublead_DecorKey);
    SG::ReadDecorHandle<xAOD::DiTauJetContainer,float> f_subjet_sublDec(m_f_subjet_subl_DecorKey);
    SG::ReadDecorHandle<xAOD::DiTauJetContainer,float> f_subjetDec(m_f_subjets_DecorKey);
    SG::ReadDecorHandle<xAOD::DiTauJetContainer,float> R_max_leadDec(m_R_max_lead_DecorKey); 
    SG::ReadDecorHandle<xAOD::DiTauJetContainer,float> R_max_subleadDec(m_R_max_sublead_DecorKey);
    SG::ReadDecorHandle<xAOD::DiTauJetContainer,int>   n_trackDec(m_n_track_DecorKey);
    SG::ReadDecorHandle<xAOD::DiTauJetContainer,float> R_isotracDec(m_R_isotrack_DecorKey);   
    SG::ReadDecorHandle<xAOD::DiTauJetContainer,float> R_tracks_sublDec(m_R_track_sublead_DecorKey);
    SG::ReadDecorHandle<xAOD::DiTauJetContainer,float> M_core_leadDec(m_M_core_lead_DecorKey);  
    SG::ReadDecorHandle<xAOD::DiTauJetContainer,float> M_core_sublDec(m_M_core_sublead_DecorKey); 
    SG::ReadDecorHandle<xAOD::DiTauJetContainer,float> M_tracks_leadDec(m_M_track_lead_DecorKey);
    SG::ReadDecorHandle<xAOD::DiTauJetContainer,float> d0_leadtrack_leadDec(m_d0_leadtrack_lead_DecorKey);
    SG::ReadDecorHandle<xAOD::DiTauJetContainer,float> d0_leadtrack_sublDec(m_d0_leadtrack_sublead_DecorKey);
    SG::ReadDecorHandle<xAOD::DiTauJetContainer,float> f_isotracks_Dec(m_f_isotracks_DecorKey);
    // Accessors for reading the necessary features from the xAOD::TrackParticle object
    static const SG::ConstAccessor< uint8_t > numberOfInrmstPxlLyrHitsAcc ("numberOfInnermostPixelLayerHits");
    static const SG::ConstAccessor< uint8_t > numberOfPixelHitsAcc        ("numberOfPixelHits");
    static const SG::ConstAccessor< uint8_t > numberOfSCTHitsAcc          ("numberOfSCTHits");
    static const SG::ConstAccessor<   float > z0Acc                       ("z0");
    static const SG::ConstAccessor<   float > d0Acc                       ("d0");
 
 
    // Input features for Ditau tagger ONNX model
    std::vector<float> jet_vars = {
        R_max_leadDec                   (ditau),
        R_max_subleadDec                (ditau),
        R_tracks_sublDec                (ditau),
        R_isotracDec                    (ditau),
        d0_leadtrack_leadDec            (ditau),
        d0_leadtrack_sublDec            (ditau),
        f_core_leadDec                  (ditau),
        f_core_sublDec                  (ditau),
        f_subjet_sublDec                (ditau),
        f_subjetDec                     (ditau),
        f_isotracks_Dec                 (ditau),
        M_core_leadDec                  (ditau),
        M_core_sublDec                  (ditau),
        M_tracks_leadDec                (ditau),
        static_cast<float>( n_trackDec  (ditau)),
    };
    std::vector<int64_t> jet_shape = {1, static_cast<int64_t>(jet_vars.size())};
 
    const TrackParticleLinks_t &vTauTracks = ditau.trackLinks();
    std::vector<std::vector<float>> track_features(m_maxTracks, std::vector<float>(11, 0.0f));
 
    float jet_eta = ditau.eta();
    float jet_phi = ditau.phi();
    size_t num_tracks = std::min(static_cast<size_t>(m_maxTracks), vTauTracks.size());
 
    for (size_t i = 0; i < num_tracks; ++i) {
        const ElementLink<xAOD::TrackParticleContainer> &trackLink = vTauTracks[i];
        if (!trackLink.isValid()) continue;
        const xAOD::TrackParticle *xTrack = *trackLink;
        float track_eta    = xTrack->eta();
        float track_phi    = xTrack->phi();
        float delta_eta    = track_eta - jet_eta;
        float delta_phi    = std::remainder(track_phi - jet_phi, 2 * M_PI);
        float delta_R      = std::hypot(delta_eta, delta_phi);
        float track_pt     = static_cast<float>(xTrack->pt());
        float pt_log       = std::log(track_pt + 1e-8f);
        float jet_pt       = ditau_ptDec(ditau);
        float pt_ratio     = track_pt / jet_pt;
        float pt_ratio_log = std::log(1.0f - pt_ratio + 1e-8f);
        float track_charge = xTrack->charge();
        float pt_ratio_log_nan_less = nan_to_num(pt_ratio_log, 0.0f, 0.0f, 0.0f);
 
        track_features[i] = {
            delta_eta,
            delta_phi,
            pt_log,
            d0Acc(*xTrack),
            pt_ratio_log_nan_less,
            z0Acc(*xTrack),
            delta_R,
            static_cast<float>(numberOfInrmstPxlLyrHitsAcc(*xTrack)),
            static_cast<float>(numberOfPixelHitsAcc(*xTrack)),
            static_cast<float>(numberOfSCTHitsAcc(*xTrack)),
            track_charge
        };
    }
    std::vector<int64_t> track_shape = {1, static_cast<int64_t>(m_maxTracks), 11};
 
    // Actual ONNX inference
    OnnxInputs inputs{
        flatten(track_features),
        track_shape,
        extract_points(track_features),
        {1, track_shape[1], 2},
        create_mask(track_features),
        {1, track_shape[1]},
        std::move(jet_vars),
        std::move(jet_shape),
        {0.0f},
        {1, 1}
    };
    auto output = run_inference(inputs);
    return output.output_1[1];
}