TauDecayModeNNClassifier.h

Go to the documentation of this file.

/*
  Copyright (C) 2002-2025 CERN for the benefit of the ATLAS collaboration
*/
 
#ifndef TAURECTOOLS_TAUDECAYMODENNCLASSIFIER_H
#define TAURECTOOLS_TAUDECAYMODENNCLASSIFIER_H
 
// base class include(s)
#include "tauRecTools/TauRecToolBase.h"
 
#include "AsgTools/PropertyWrapper.h"
 
// xAOD include(s)
#include "xAODTau/TauJet.h"
 
// lwtnn include(s)
#include "lwtnn/LightweightGraph.hh"
#include "lwtnn/parse_json.hh"
#include "lwtnn/Exceptions.hh"
 
// standard library include(s)
#include <memory>
#include <vector>
#include <set>
#include <map>

 
class TauDecayModeNNClassifier : public TauRecToolBase
{
public:
  ASG_TOOL_CLASS2(TauDecayModeNNClassifier, TauRecToolBase, ITauToolBase)
 
  explicit TauDecayModeNNClassifier(const std::string &name = "TauDecayModeNNClassifier");
  virtual ~TauDecayModeNNClassifier();
 
  virtual StatusCode initialize() override;
  virtual StatusCode execute(xAOD::TauJet &xTau) const override;
 
private:
  Gaudi::Property<std::string> m_outputName{this, "OutputName", "NNDecayMode"};
  Gaudi::Property<std::string> m_probPrefix{this, "ProbPrefix", "NNDecayModeProb_"};
  Gaudi::Property<std::string> m_weightFile{this, "WeightFile", ""};
  Gaudi::Property<std::size_t> m_maxTauTracks{this, "MaxTauTracks", 3};
  Gaudi::Property<std::size_t> m_maxNeutralPFOs{this, "MaxNeutralPFOs", 8};
  Gaudi::Property<std::size_t> m_maxShotPFOs{this, "MaxShotPFOs", 6};
  Gaudi::Property<std::size_t> m_maxConvTracks{this, "MaxConvTracks", 4};
  Gaudi::Property<float> m_neutralPFOPtCut{this, "NeutralPFOPtCut", 1.5};
  Gaudi::Property<bool> m_ensureTrackConsistency{this, "EnsureTrackConsistency", true};
  Gaudi::Property<bool> m_decorateProb{this, "DecorateProb", true};

  virtual StatusCode getInputs(const xAOD::TauJet &xTau,
                               std::map<std::string, std::map<std::string, std::vector<double>>> &inputSeqMap) const;
  std::unique_ptr<const lwt::LightweightGraph> m_lwtGraph; 
};
 
namespace tauRecTools
{
  class TauDecayModeNNVariable
  {
  public:
    TauDecayModeNNVariable() = delete;
    static const std::size_t nClasses = 5;
    static const std::set<std::string> sCommonP4Vars;
    static const std::set<std::string> sTrackIPVars;
    static const std::set<std::string> sNeutralPFOVars;
    static const std::array<std::string, nClasses> sModeNames;
    static float deltaPhi(const TLorentzVector &p4, const TLorentzVector &p4_tau);
    static float deltaEta(const TLorentzVector &p4, const TLorentzVector &p4_tau);
    static float deltaPhiECal(const TLorentzVector &p4, const std::pair<float, bool> &tau_phiTrkECal);
    static float deltaEtaECal(const TLorentzVector &p4, const std::pair<float, bool> &tau_etaTrkECal);
    template <typename T>
    static T pfoAttr(const xAOD::PFO *pfo, const xAOD::PFODetails::PFOAttributes &attr);
  };

  class TauDecayModeNNHelper
  {
  public:
    TauDecayModeNNHelper() = delete;
    static float Log10Robust(const float val, const float min_val = 0.);
    template <typename T>
    static void sortAndKeep(std::vector<T> &vec, const std::size_t n_obj);
    template <typename T>
    static void initMapKeys(std::map<std::string, T> &empty_map, const std::set<std::string> &keys);
  };
} // namespace tauRecTools
 
#endif // TAURECTOOLS_TAUDECAYMODENNCLASSIFIER_H