JetTruthLabelingTool.h

Go to the documentation of this file.

/*
  Copyright (C) 2002-2026 CERN for the benefit of the ATLAS collaboration
*/
 
#ifndef PARTICLEJETTOOLS_JETTRUTHLABELINGTOOL_H
#define PARTICLEJETTOOLS_JETTRUTHLABELINGTOOL_H
 
#include "AsgTools/AsgTool.h"
#include "AsgTools/PropertyWrapper.h"
#include "AsgDataHandles/ReadHandleKey.h"
#include "AsgDataHandles/ReadDecorHandleKey.h"
#include "AsgDataHandles/ReadDecorHandle.h"
#include "AsgDataHandles/WriteDecorHandleKey.h"
#include "AsgDataHandles/WriteDecorHandle.h"
 
#include "xAODJet/JetContainer.h"
#include "xAODTruth/TruthParticleContainer.h"
#include "xAODEventInfo/EventInfo.h"
#include "AthLinks/ElementLink.h"
#include "JetInterface/IJetDecorator.h"
#include "ParticleJetTools/LargeRJetLabelEnum.h"
#include "ParticleJetTools/SmallRJetPileupLabelEnum.h"
#include <optional>
 
class JetTruthLabelingTool :
  public asg::AsgTool,
  virtual public IJetDecorator
{
  ASG_TOOL_CLASS(JetTruthLabelingTool, IJetDecorator)
 
public:

  JetTruthLabelingTool(const std::string& name = "JetTruthLabelingTool");
  virtual StatusCode initialize() override;

  StatusCode decorate(const xAOD::JetContainer& jets) const override;

  void print() const override;
  
  const std::string& getLargeRJetTruthLabelName() const {
    return m_truthLabelName;
  };
 
protected:
  
  Gaudi::Property<std::string> m_jetContainerName{this, "RecoJetContainer", "", "Input reco jet container name"};
  Gaudi::Property<std::string> m_truthLabelName{this, "TruthLabelName", "R10TruthLabel_R22v1", "Truth label name"};
 
  Gaudi::Property<bool> m_isTruthJetCol{this, "IsTruthJetCollection", false, "Flag indicating whether input collection is a truth jet container"};
  Gaudi::Property<bool> m_forceDeltaRMatch{this, "ForceDeltaRMatch", false, "Whether to force dR matching e.g. if GhostTruth is not available"};
 
  SG::ReadHandleKey<xAOD::EventInfo> m_evtInfoKey{this, "EventInfoKey", "EventInfo", "Name of EventInfo object"};

  SG::ReadHandleKey<xAOD::JetContainer> m_OOTJetName{ this, "OOTJetContainer", "OutOfTimeAntiKt4TruthJets", "Do not configure manually!" };
  SG::ReadHandleKey<xAOD::JetContainer> m_ITJetName{ this, "ITJetContainer", "InTimeAntiKt4TruthJets", "Do not configure manually!" };
  SG::ReadHandleKey<xAOD::JetContainer> m_truthJetCollectionKey{this, "TruthJetContainer", "", "Do not configure manually!"};
  SG::ReadHandleKey<xAOD::JetContainer> m_truthGroomedJetCollectionKey{this, "TruthGroomedJetContainer", "", "Do not configure manually!"};
  bool m_useGhostJetMatch{}; 
  bool m_isSmallR{}; 
  bool m_matchUngroomedParent{}; 
  bool m_getTruthGroomedJetValues{}; 
  double m_dRTruthJet{}; 
  double m_recoGhostFrac{}; 
  double m_mLowTop{}; 
  double m_mLowW{}; 
  double m_mLowZ{}; 
  bool m_doLargeRLabels{true}; 
 
  struct DecorHandles {
    DecorHandles (const JetTruthLabelingTool& tool, const EventContext& ctx);
 
    using IntHandle_t = SG::WriteDecorHandle<xAOD::JetContainer, int>;
    using IntHandleOp_t = std::optional<IntHandle_t>;
    IntHandleOp_t labelHandle;
    IntHandleOp_t nbHandle;
    IntHandleOp_t labelRecoHandle;
    IntHandleOp_t nbRecoHandle;
    IntHandleOp_t matchedPileupTagHandle;
 
    using FloatHandle_t = SG::WriteDecorHandle<xAOD::JetContainer, float>;
    using FloatHandleOp_t = std::optional<FloatHandle_t>;
    FloatHandleOp_t split23Handle;
    FloatHandleOp_t split12Handle;
    using ELHandle_t = SG::WriteDecorHandle<xAOD::JetContainer, ElementLink<xAOD::JetContainer> >;
    using ELHandleOp_t = std::optional<ELHandle_t>;
    ELHandleOp_t matchedTruthJetHandle;
    FloatHandleOp_t matchedTruthJetMassHandle;
    FloatHandleOp_t matchedTruthJetPtHandle;
    FloatHandleOp_t matchedTruthJetRapidityHandle;
    FloatHandleOp_t matchedTruthJetPhiHandle;
    FloatHandleOp_t matchedTruthJetEHandle;
    FloatHandleOp_t matchedTruthJetDRHandle;
    FloatHandleOp_t matchedTruthJetGFHandle;
    FloatHandleOp_t sumPtMatchedHSJetsHandle;
    FloatHandleOp_t sumPtMatchedOOTJetsHandle;
    FloatHandleOp_t sumPtMatchedITJetsHandle;
    FloatHandleOp_t matchedTruthGroomedMassHandle;
    FloatHandleOp_t matchedTruthGroomedPtHandle;
  };
  friend struct DecorHandles;

  StatusCode labelTruthJets( DecorHandles& dh,
                             const EventContext& ctx ) const;
  StatusCode labelTruthJets( DecorHandles& dh,
                             const xAOD::JetContainer &jets,
                             const EventContext& ctx ) const;

  StatusCode labelRecoJets(DecorHandles& dh,
                           const xAOD::JetContainer &jets,
                           const EventContext& ctx) const;

  int getTruthJetLabelGA( DecorHandles& dh,
                          const xAOD::Jet &jet,
                          const EventContext& ctx ) const;

  int getLabel( DecorHandles& dh,
                const xAOD::Jet &jet, bool matchH, bool matchW, bool matchZ, bool matchTop,
                const EventContext& ctx ) const;

  float getWZSplit12Cut( float pt ) const;

  float getTopSplit23Cut( float pt ) const;

  int getNGhostParticles( const xAOD::Jet &jet, const std::string & collection ) const;
 
  enum class TruthLabelConfiguration {
      R21Precision_2022v1,
      R10TruthLabel_R22v1,
      R10WZTruthLabel_R22v1,
      R4TruthLabel,
      R4TruthDressedWZLabel,
      R4InTimeTruthLabel,
      R4OutOfTimeTruthLabel,
      Unknown
  };
 
  TruthLabelConfiguration parseLabel(const std::string& label);
  TruthLabelConfiguration m_truthLabelConfig = TruthLabelConfiguration::Unknown;
 
  SG::WriteDecorHandleKey<xAOD::JetContainer> m_label_truthKey{this, "label_TruthKey", "", "Do not configure manually!"};
  SG::WriteDecorHandleKey<xAOD::JetContainer> m_NB_truthKey{this, "NB_TruthKey", "", "Do not configure manually!"};
  SG::ReadDecorHandleKey<xAOD::JetContainer> m_split12_truthKey{this, "Split12_TruthKey", "", "Do not configure manually!"};
  SG::ReadDecorHandleKey<xAOD::JetContainer> m_split23_truthKey{this, "Split23_TruthKey", "", "Do not configure manually!"};
 
  SG::WriteDecorHandleKey<xAOD::JetContainer> m_label_recoKey{this, "label_RecoKey", "", "Do not configure manually!"};
  SG::WriteDecorHandleKey<xAOD::JetContainer> m_NB_recoKey{this, "NB_RecoKey", "", "Do not configure manually!"};
  SG::WriteDecorHandleKey<xAOD::JetContainer> m_truthSplit12_recoKey{this, "TruthSplit12_RecoKey", "", "Do not configure manually!"};
  SG::WriteDecorHandleKey<xAOD::JetContainer> m_truthSplit23_recoKey{this, "TruthSplit23_RecoKey", "", "Do not configure manually!"};
 
  SG::WriteDecorHandleKey<xAOD::JetContainer> m_matchedTruthJet_recoKey{this, "TruthMatch_Jet_RecoKey", "", "Do not configure manually!"};
  SG::WriteDecorHandleKey<xAOD::JetContainer> m_matchedTruthJetMass_recoKey{this, "TruthMatch_m_RecoKey", "", "Do not configure manually!"};
  SG::WriteDecorHandleKey<xAOD::JetContainer> m_matchedTruthJetPt_recoKey{this, "TruthMatch_pt_RecoKey", "", "Do not configure manually!"};
  SG::WriteDecorHandleKey<xAOD::JetContainer> m_matchedTruthJetRapidity_recoKey{ this, "TruthMatch_rapidity_RecoKey", "", "Do not configure manually!" };
  SG::WriteDecorHandleKey<xAOD::JetContainer> m_matchedTruthJetPhi_recoKey{ this, "TruthMatch_phi_RecoKey", "", "Do not configure manually!" };
  SG::WriteDecorHandleKey<xAOD::JetContainer> m_matchedTruthJetDR_recoKey{ this, "TruthMatch_dR_RecoKey", "", "Do not configure manually!" };
  SG::WriteDecorHandleKey<xAOD::JetContainer> m_matchedPileupTag_recoKey{ this, "TruthMatch_PileupLabel_RecoKey", "", "Do not configure manually!" };
  SG::WriteDecorHandleKey<xAOD::JetContainer> m_sumPtMatchedHSJets_recoKey{ this, "TruthMatch_HSSumPt_RecoKey", "", "Do not configure manually!" };
  SG::WriteDecorHandleKey<xAOD::JetContainer> m_sumPtMatchedOOTJets_recoKey{ this, "TruthMatch_OOTSumPt_RecoKey", "", "Do not configure manually!" };
  SG::WriteDecorHandleKey<xAOD::JetContainer> m_sumPtMatchedITJets_recoKey{ this, "TruthMatch_ITSumPt_RecoKey", "", "Do not configure manually!" };
 
  SG::WriteDecorHandleKey<xAOD::JetContainer> m_matchedTruthGroomedJetMass_recoKey{this, "TruthMatch_GroomedJetMass_RecoKey", "", "Do not configure manually!"};
  SG::WriteDecorHandleKey<xAOD::JetContainer> m_matchedTruthGroomedJetPt_recoKey{this, "TruthMatch_GroomedJetPt_RecoKey", "", "Do not configure manually!"};
};
 
#endif