JetTruthLabelingTool.h

Go to the documentation of this file.

/*
  Copyright (C) 2002-2024 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 "JetInterface/IJetDecorator.h"
#include "ParticleJetTools/LargeRJetLabelEnum.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_R21Consolidated", "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"};

  Gaudi::Property<bool> m_useTRUTH3{this, "UseTRUTH3", false, "True for TRUTH3 format, false for TRUTH1"};
  SG::ReadHandleKey<xAOD::TruthParticleContainer> m_truthParticleContainerName{this, "TruthParticleContainerName", "TruthParticles", "Truth particle container name"};
  SG::ReadHandleKey<xAOD::TruthParticleContainer> m_truthBosonContainerName{this, "TruthBosonContainerName", "TruthBosonsWithDecayParticles", "Truth boson container name"};
  SG::ReadHandleKey<xAOD::TruthParticleContainer> m_truthTopQuarkContainerName{this, "TruthTopQuarkContainerName", "TruthTopQuarkWithDecayParticles", "Truth top container name"};

  SG::ReadHandleKey<xAOD::JetContainer> m_truthJetCollectionName{this, "TruthJetContainer", "", "Do not configure manually!"};
  SG::ReadHandleKey<xAOD::JetContainer> m_truthGroomedJetCollectionName{this, "TruthGroomedJetContainer", "", "Do not configure manually!"};
  bool m_useDRMatch{}; 
  bool m_useWZMassHigh{}; 
  bool m_useGhostJetMatch{}; 
  bool m_matchUngroomedParent{}; 
  bool m_getTruthGroomedJetValues{}; 
  double m_dRTruthJet{}; 
  double m_recoGhostFrac{}; 
  double m_dRTruthPart{}; 
  double m_mLowTop{}; 
  double m_mLowW{}; 
  double m_mHighW{}; 
  double m_mLowZ{}; 
  double m_mHighZ{}; 
 
  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;
    using FloatHandle_t = SG::WriteDecorHandle<xAOD::JetContainer, float>;
    using FloatHandleOp_t = std::optional<FloatHandle_t>;
    FloatHandleOp_t dRWHandle;
    FloatHandleOp_t dRZHandle;
    FloatHandleOp_t dRHHandle;
    FloatHandleOp_t dRTopHandle;
    FloatHandleOp_t dRWRecoHandle;
    FloatHandleOp_t dRZRecoHandle;
    FloatHandleOp_t dRHRecoHandle;
    FloatHandleOp_t dRTopRecoHandle;
    FloatHandleOp_t split23Handle;
    FloatHandleOp_t split12Handle;
    FloatHandleOp_t matchedTruthJetMassHandle;
    FloatHandleOp_t matchedTruthJetPtHandle;
    FloatHandleOp_t matchedTruthJetEtaHandle;
    FloatHandleOp_t matchedTruthJetPhiHandle;
    FloatHandleOp_t matchedTruthJetDRHandle;
    FloatHandleOp_t matchedTruthJetGFHandle;
    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 getTruthJetLabelDR( DecorHandles& dh,
                          const xAOD::Jet &jet,
                          const std::vector<std::pair<TLorentzVector,int> >& tlv_truthParts,
                          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, std::string collection ) const;

  bool getIsSherpa( const int DSID ) const {
    if(
        (304307 <= DSID && DSID <= 304309) || // Sherpa 2.2.1 W+jets
        (304707 <= DSID && DSID <= 304709) // Sherpa 2.2.1 Z+jets
    ) { 
      return true;
    }
    return false;
  };

  void getTLVs( std::vector< std::pair<TLorentzVector,int> > &tlvs, const xAOD::TruthParticleContainer *truthBosons, const xAOD::TruthParticleContainer *truthTop, bool isSherpa = false ) const;

  bool selectTruthParticle( const xAOD::TruthParticle *tp, int pdgId ) const;
 
  enum class TruthLabelConfiguration {
      R21Consolidated,
      R21Precision,
      R21Precision_2022v1,
      R10TruthLabel_R22v1,
      R10WZTruthLabel_R22v1,
      R4TruthLabel,
      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_dR_W_truthKey{this, "dR_W_TruthKey", "", "Do not configure manually!"};
  SG::WriteDecorHandleKey<xAOD::JetContainer> m_dR_Z_truthKey{this, "dR_Z_TruthKey", "", "Do not configure manually!"};
  SG::WriteDecorHandleKey<xAOD::JetContainer> m_dR_H_truthKey{this, "dR_H_TruthKey", "", "Do not configure manually!"};
  SG::WriteDecorHandleKey<xAOD::JetContainer> m_dR_Top_truthKey{this, "dR_Top_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_dR_W_recoKey{this, "dR_W_RecoKey", "", "Do not configure manually!"};
  SG::WriteDecorHandleKey<xAOD::JetContainer> m_dR_Z_recoKey{this, "dR_Z_RecoKey", "", "Do not configure manually!"};
  SG::WriteDecorHandleKey<xAOD::JetContainer> m_dR_H_recoKey{this, "dR_H_RecoKey", "", "Do not configure manually!"};
  SG::WriteDecorHandleKey<xAOD::JetContainer> m_dR_Top_recoKey{this, "dR_Top_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_matchedTruthJetMass_recoKey{this, "MatchedTruthJetMass_RecoKey", "", "Do not configure manually!"};
  SG::WriteDecorHandleKey<xAOD::JetContainer> m_matchedTruthJetPt_recoKey{this, "MatchedTruthJetPt_RecoKey", "", "Do not configure manually!"};
  SG::WriteDecorHandleKey<xAOD::JetContainer> m_matchedTruthJetEta_recoKey{ this, "MatchedTruthJetEta_RecoKey", "", "Do not configure manually!" };
  SG::WriteDecorHandleKey<xAOD::JetContainer> m_matchedTruthJetPhi_recoKey{ this, "MatchedTruthJetPhi_RecoKey", "", "Do not configure manually!" };
  SG::WriteDecorHandleKey<xAOD::JetContainer> m_matchedTruthJetDR_recoKey{ this, "MatchedTruthJetDR_RecoKey", "", "Do not configure manually!" };
  SG::WriteDecorHandleKey<xAOD::JetContainer> m_matchedTruthJetGF_recoKey{ this, "MatchedTruthJetGF_RecoKey", "", "Do not configure manually!" };
 
  SG::WriteDecorHandleKey<xAOD::JetContainer> m_matchedTruthGroomedJetMass_recoKey{this, "MatchedTruthGroomedJetMass_RecoKey", "", "Do not configure manually!"};
  SG::WriteDecorHandleKey<xAOD::JetContainer> m_matchedTruthGroomedJetPt_recoKey{this, "MatchedTruthGroomedJetPt_RecoKey", "", "Do not configure manually!"};
};
 
#endif