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>
 
#include <map>
 
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"};
 
  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!"};
 
  Gaudi::Property<bool> m_useGhostJetMatch{this, "UseGhostJetMatch", false, "Use ghost association to match reco to truth jets, dR otherwise"};
  Gaudi::Property<double> m_recoGhostFrac{this, "RecoGhostFraction", 0.75, "Ghost pT fraction window for GA matching"};
  Gaudi::Property<bool> m_matchUngroomedParent{this, "MatchUngroomedParent", false, "Use the ungroomed reco jet parent to match to truth jet"};
  Gaudi::Property<bool> m_getTruthGroomedJetValues{this, "GetTruthGroomedJetValues", false, "Save properties of the groomed truth jet"};
  Gaudi::Property<double> m_dRTruthJet{this, "dRTruthJet", 0.0, "dR to match truth jet to reco jet (<0: closest jet)"};
  Gaudi::Property<bool> m_doLargeRLabels{this, "DoLargeRLabels", true, "Compute the W/Z/top labels (large R jets)"};

  Gaudi::Property<double> m_mLowTop{this, "MassLowTop", 140.0, "Lower mass cut for top label [GeV]"};
  Gaudi::Property<double> m_mLowW{this, "MassLowW", 50.0, "Lower mass cut for W label [GeV]"};
  Gaudi::Property<double> m_mLowZ{this, "MassLowZ", 50.0, "Lower mass cut for Z label [GeV]"};
 
  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