TauGNNEvaluator.cxx

Go to the documentation of this file.

/*
  Copyright (C) 2002-2025 CERN for the benefit of the ATLAS collaboration
*/
 
#include "tauRecTools/TauGNNEvaluator.h"
#include "tauRecTools/HelperFunctions.h"
 
#include "PathResolver/PathResolver.h"
 
#include <algorithm>
 
 
TauGNNEvaluator::TauGNNEvaluator(const std::string &name): 
  TauRecToolBase(name),
  m_net_inclusive(nullptr),
  m_net_0p(nullptr), m_net_1p(nullptr), m_net_2p(nullptr), m_net_3p(nullptr) {
}
 
TauGNNEvaluator::~TauGNNEvaluator() {}
 
StatusCode TauGNNEvaluator::initialize() {
  ATH_MSG_INFO("Initializing TauGNNEvaluator with "<<m_max_tracks.value()<<" tracks and "<<m_max_clusters<<" clusters...");
 
  // Set the layer and node names in the weight file
  TauGNN::Config config;
  config.input_layer_scalar = m_input_layer_scalar;
  config.input_layer_tracks = m_input_layer_tracks;
  config.input_layer_clusters = m_input_layer_clusters;
  config.output_node_tau = m_outnode_tau;
  config.output_node_jet = m_outnode_jet;
 
  // We can either use an inclussive GNN (e.g. Offline GNTauv0), or a prong-dependent GNN (e.g. HLT GNTau), not both!
  
  if(!m_weightfile_inclusive.empty()) { // Prong-inclusive network
    if(!m_weightfile_0p.empty() || !m_weightfile_1p.empty() || !m_weightfile_2p.empty() || !m_weightfile_3p.empty()) {
      ATH_MSG_ERROR("Cannot load both prong-inclusive and prong-dependent networks!");
      return StatusCode::FAILURE;
    }
    
    ATH_MSG_INFO("Loading prong-inclusive TauID GNN");
    m_net_inclusive = load_network(m_weightfile_inclusive, config);
    if(!m_net_inclusive) return StatusCode::FAILURE;
 
  } else { // Prong-dependent networks
 
    // 0-prong is optional
    if(!m_weightfile_0p.empty()) {
      ATH_MSG_INFO("Loading 0-prong TauID GNN");
      m_net_0p = load_network(m_weightfile_0p, config);
      if(!m_net_0p) return StatusCode::FAILURE;
    }
 
    ATH_MSG_INFO("Loading 1-prong TauID GNN");
    m_net_1p = load_network(m_weightfile_1p, config);
    if(!m_net_1p) return StatusCode::FAILURE;
 
    // 2-prong is optional
    if(!m_weightfile_2p.empty()) {
      ATH_MSG_INFO("Loading 2-prong TauID GNN");
      m_net_2p = load_network(m_weightfile_2p, config);
      if(!m_net_2p) return StatusCode::FAILURE;
    }
 
    ATH_MSG_INFO("Loading 3-prong TauID GNN");
    m_net_3p = load_network(m_weightfile_3p, config);
    if(!m_net_3p) return StatusCode::FAILURE;
  }
 
  if(m_output_discriminant < Discriminant::NegLogPJet || m_output_discriminant > Discriminant::PTau) {
    ATH_MSG_FATAL("Invalid TauGNNEvaluator discriminant setting: " << m_output_discriminant);
  }
 
  if (!m_tauContainerName.empty()){
    m_scoreHandleKey = m_tauContainerName + "." + m_output_varname;
    ATH_CHECK(m_scoreHandleKey.initialize());    
  }
 
  return StatusCode::SUCCESS;
}
 
std::unique_ptr<TauGNN> TauGNNEvaluator::load_network(const std::string& network_file, const TauGNN::Config& config) const {
  // Use PathResolver to search for the weight files
  if(network_file.empty()) return nullptr;
 
  const std::string pr_network_file = find_file(network_file);
  if(pr_network_file.empty()) {
    ATH_MSG_ERROR("Could not find network weights: " << network_file);
    return nullptr;
  }
  
  ATH_MSG_INFO("Using network config: " << pr_network_file);
  
  // Load the weights and create the network
  std::unique_ptr<TauGNN> net = std::make_unique<TauGNN>(pr_network_file, config, m_useTRT);
  if(!net) ATH_MSG_ERROR("No network configured.");
 
  return net;
}
 
StatusCode TauGNNEvaluator::execute(xAOD::TauJet &tau) const {
  // Output variable Decorators
  const SG::Accessor<float> output(m_output_varname);
  const SG::Accessor<float> out_ptau(m_output_ptau);
  const SG::Accessor<float> out_pjet(m_output_pjet);
  const SG::Decorator<char> out_trkclass("GNTau_TrackClass");
  // Set default score and overwrite later
  output(tau) = -1111.0f;
  out_ptau(tau) = -1111.0f;
  out_pjet(tau) = -1111.0f;
 
  //Skip execution for low-pT taus to save resources
  if (tau.pt() < m_minTauPt) {
    return StatusCode::SUCCESS;
  }
 
  // save CPU when running PHYS derivations
  if (m_applyLooseTrackSel) {
    if (tau.nTracks()>5) return StatusCode::SUCCESS;
  }
 
  // save CPU when running in RAWtoALL for tau trigger monitoring purpose
  if (m_applyTightTrackSel) {
    if (tau.nTracks()!=1 && tau.nTracks()!=3) return StatusCode::SUCCESS;
  }
 
  // Get input objects
  ATH_MSG_DEBUG("Fetching Tracks");
  std::vector<const xAOD::TauTrack *> tracks;
  ATH_CHECK(get_tracks(tau, tracks));
  ATH_MSG_DEBUG("Fetching clusters");
  std::vector<xAOD::CaloVertexedTopoCluster> clusters;
  ATH_CHECK(get_clusters(tau, clusters));
  ATH_MSG_DEBUG("Constituent fetching done...");
 
  // Truncate tracks
  int numTracksMax = std::min(m_max_tracks.value(), static_cast<int>(tracks.size()));
  std::vector<const xAOD::TauTrack *> trackVec(tracks.begin(), tracks.begin()+numTracksMax);
 
  // Network outputs
  std::map<std::string, float> out_f;
  std::map<std::string, std::vector<char>> out_vc;
  std::map<std::string, std::vector<float>> out_vf;
 
  // Evaluate networks
  if(m_net_inclusive) {
    std::tie(out_f, out_vc, out_vf) = m_net_inclusive->compute(tau, trackVec, clusters);
  } else {
    // First we calculate the tau prongness
    int n_tracks = tau.nTracksCharged();
    if(m_min_prong_track_pt) {
      n_tracks = 0;
      for(const xAOD::TauTrack* track : tracks) {
        if(track->pt() > m_min_prong_track_pt) n_tracks++;
      }
    }
    ATH_MSG_DEBUG("Tau prongness: " << n_tracks);
 
    if(n_tracks == 0 && m_net_0p) std::tie(out_f, out_vc, out_vf) = m_net_0p->compute(tau, trackVec, clusters);
    else if(n_tracks == 1) std::tie(out_f, out_vc, out_vf) = m_net_1p->compute(tau, trackVec, clusters);
    else if(n_tracks == 2) {
      if(m_net_2p) std::tie(out_f, out_vc, out_vf) = m_net_2p->compute(tau, trackVec, clusters);
      else std::tie(out_f, out_vc, out_vf) = m_net_3p->compute(tau, trackVec, clusters);
    } else if(n_tracks == 3) std::tie(out_f, out_vc, out_vf) = m_net_3p->compute(tau, trackVec, clusters);
  }
 
  // Store scores only if the inferences actually ran
  if(out_f.contains(m_outnode_tau)) {
    if(m_output_discriminant == Discriminant::NegLogPJet) {
        output(tau) = std::log10(1/(1-out_f.at(m_outnode_tau)));
    } else if(m_output_discriminant == Discriminant::PTau) {
        output(tau) = out_f.at(m_outnode_tau);
    }
 
    out_ptau(tau) = out_f.at(m_outnode_tau);
    out_pjet(tau) = out_f.at(m_outnode_jet);
  }
 
  return StatusCode::SUCCESS;
}
 
 
StatusCode TauGNNEvaluator::get_tracks(const xAOD::TauJet &tau, std::vector<const xAOD::TauTrack *> &out) const {
  std::vector<const xAOD::TauTrack*> tracks = tau.allTracks();
 
  // Skip unclassified tracks:
  // - the track is a LRT and classifyLRT = false
  // - the track is not among the MaxNtracks highest-pt tracks in the track classifier
  // - track classification is not run (trigger)
  if(m_doTrackClassification) {
    std::vector<const xAOD::TauTrack*>::iterator it = tracks.begin();
    while(it != tracks.end()) {
      if((*it)->flag(xAOD::TauJetParameters::unclassified)) {
    it = tracks.erase(it);
      }
      else {
    ++it;
      }
    }
  }
 
  // Sort by descending pt
  auto cmp_pt = [](const xAOD::TauTrack *lhs, const xAOD::TauTrack *rhs) {
    return lhs->pt() > rhs->pt();
  };
  std::sort(tracks.begin(), tracks.end(), cmp_pt);
  out = std::move(tracks);
 
  return StatusCode::SUCCESS;
}
 
StatusCode TauGNNEvaluator::get_clusters(const xAOD::TauJet &tau, std::vector<xAOD::CaloVertexedTopoCluster> &clusters) const {
 
  TLorentzVector tauAxis = tauRecTools::getTauAxis(tau, m_doVertexCorrection);
 
  for (const xAOD::CaloVertexedTopoCluster& vertexedCluster : tau.vertexedClusters()) {
    TLorentzVector clusterP4 = vertexedCluster.p4();
    if (clusterP4.DeltaR(tauAxis) > m_max_cluster_dr) continue;
      
    clusters.push_back(vertexedCluster);
  }
 
  // Sort by descending et
  auto et_cmp = [](const xAOD::CaloVertexedTopoCluster& lhs,
           const xAOD::CaloVertexedTopoCluster& rhs) {
    return lhs.p4().Et() > rhs.p4().Et();
  };
  std::sort(clusters.begin(), clusters.end(), et_cmp);
 
  // Truncate clusters
  if (static_cast<int>(clusters.size()) > m_max_clusters) {
    clusters.resize(m_max_clusters, clusters[0]);
  }
 
  return StatusCode::SUCCESS;
}