TausRUsEvaluator.h

Go to the documentation of this file.

/*
  Copyright (C) 2002-2026 CERN for the benefit of the ATLAS collaboration
*/
 
#ifndef TAURECTOOLS_TAUSRUSEVALUATOR_H
#define TAURECTOOLS_TAUSRUSEVALUATOR_H
 
#include "tauRecTools/TauRecToolBase.h"
#include "tauRecTools/TausRUsDataLoader.h"
 
#include "AthOnnxInterfaces/IAthInferenceTool.h"
 
#include "AsgDataHandles/ReadHandleKey.h"
#include "AsgDataHandles/WriteDecorHandleKey.h"
#include "AsgTools/PropertyWrapper.h"
#include "AsgTools/ToolHandle.h"
 
#include "AthContainers/Accessor.h"
#include "AthContainers/Decorator.h"
 
#include "xAODTau/TauJet.h"
#include "xAODTau/TauJetContainer.h"
#include "xAODTau/TauTrack.h"
#include "xAODTau/TauTrackContainer.h"
#include "xAODTracking/Vertex.h"
#include "xAODTracking/VertexContainer.h"
 
#include <cstddef>
#include <cstdint>
#include <memory>
#include <span>
#include <string>
#include <vector>

class TausRUsEvaluator : public TauRecToolBase {
public:
  ASG_TOOL_CLASS2(TausRUsEvaluator, TauRecToolBase, ITauToolBase)
 
  TausRUsEvaluator(const std::string& name = "TausRUsEvaluator");
  virtual ~TausRUsEvaluator();
 
  virtual StatusCode initialize() override;
  virtual StatusCode execute(xAOD::TauJet& tau) const override;

  static constexpr float DEFAULT_VALUE = -1111.0f;
  static constexpr int DEFAULT_CLASS = -1;
 
  struct Output {
    std::string name;                   
    std::string type;                   
    std::vector<int64_t> dims;
    size_t size{0};                     
  };

  enum PrimaryClass : size_t { QCD = 0, Tau = 1, Electron = 2 };

  static constexpr int TAU_TRACK_CLASS = 0;

  enum PhiComponent : size_t { PHI_SIN = 0, PHI_COS = 1 };
 
private:
 
  using FourMomDecorators = std::vector<SG::Accessor<float>>;

  StatusCode readModel(const std::string& path);
 
  void setDefaults(xAOD::TauJet& tau) const;

  void decorateFourMomentum(xAOD::TauJet& tau, const FourMomDecorators& decorators,
                            std::span<const float> ptQuantiles,
                            std::span<const float> etaValues,
                            std::span<const float> phiValues,
                            float mass) const;

  void decorateVertex(xAOD::TauJet& tau,
                      std::span<const float> scores,
                      const std::vector<const xAOD::Vertex*>& vertices) const;

  void decorateTracks(xAOD::TauJet& tau, std::span<const float> scores) const;
 
  std::vector<std::string> tauDecorationNames() const;
  std::vector<std::string> trackDecorationNames() const;
 
  ToolHandle<AthInfer::IAthInferenceTool> m_inferenceTool{
    this, "InferenceTool", "", "ONNX Runtime or Triton inference backend"};
 
  SG::ReadHandleKey<xAOD::VertexContainer> m_vertexInputContainer{
    this, "Key_vertexInputContainer", "PrimaryVertices", "Input vertex container key"};
 
  Gaudi::Property<std::string> m_modelFile{
    this, "ModelFile", "",
    "ONNX model the input and output nodes are read from, also when inference runs on Triton"};
  Gaudi::Property<std::string> m_tauContainerName{
    this, "TauContainerName", "",
    "Name of the TauJetContainer, needed to declare the output data dependencies"};
  Gaudi::Property<std::string> m_tauTrackContainerName{
    this, "TauTrackContainerName", "",
    "Name of the TauTrackContainer, needed to declare the per-track output data dependencies"};
  Gaudi::Property<float> m_minTauPt{
    this, "MinTauPt", 15., "Skip taus below this pt to save CPU"};
 
  std::unique_ptr<TausRUsDataLoader> m_loader;

  std::vector<Output> m_outputs;

  size_t m_nDecayModes{0};
  size_t m_nTrackClasses{0};
  size_t m_nQuantiles{0};
  size_t m_ptMedianQuantile{0};

  SG::Accessor<float> m_tauIDScore;
  SG::Accessor<float> m_eleRejScore;
  SG::Accessor<float> m_decayMode;
  SG::Accessor<float> m_tauCharge;
  std::vector<SG::Accessor<float>> m_decayModeScores;
  FourMomDecorators m_tauP4;
  FourMomDecorators m_chargedPionP4;
  FourMomDecorators m_neutralPionP4;
  std::vector<SG::Accessor<float>> m_vertexPosition;
  // decorators to the tracks
  SG::Decorator<int> m_trackClass;
  std::vector<SG::Decorator<float>> m_trackScores;
 
  std::vector<SG::WriteDecorHandleKey<xAOD::TauJetContainer>> m_decorKeys;
  std::vector<SG::WriteDecorHandleKey<xAOD::TauTrackContainer>> m_trackDecorKeys;
};
 
#endif // TAURECTOOLS_TAUSRUSEVALUATOR_H