ParticleJetToolsConfig.py

Go to the documentation of this file.

# Copyright (C) 2002-2025 CERN for the benefit of the ATLAS collaboration
 

 
from AthenaCommon import Logging
jrtlog = Logging.logging.getLogger('ParticleJetToolsConfig')
 
from AthenaConfiguration.ComponentFactory import CompFactory
# workaround for missing JetRecConfig in AthAnalysis
try:
    from JetRecConfig.JetRecConfig import isAnalysisRelease
except ModuleNotFoundError:
    def isAnalysisRelease():
        return True
 
# Putting MCTruthClassifier here as we needn't stick jet configs in really foreign packages
def getMCTruthClassifier():
    # Assume mc15 value
    truthclassif = CompFactory.MCTruthClassifier(
        "JetMCTruthClassifier"
        )
    if not isAnalysisRelease() :
        truthclassif.xAODTruthLinkVector= ""
    # Config neessary only for Athena releases
    import os
    if "AtlasProject" in os.environ.keys():
        if os.environ["AtlasProject"] in ["Athena","AthDerivation"]:
            truthclassif.ParticleCaloExtensionTool=""
    return truthclassif
 
# Generates truth particle containers for truth labeling
truthpartoptions = {
    "Partons":{"ToolType":CompFactory.CopyTruthPartons,"ptmin":5000},
    "BosonTop":{"ToolType":CompFactory.CopyBosonTopLabelTruthParticles,"ptmin":100000},
    "FlavourLabel":{"ToolType":CompFactory.CopyFlavorLabelTruthParticles,"ptmin":5000},
}
def getCopyTruthLabelParticles(truthtype):
    toolProperties = {}
    if truthtype == "Partons":
        truthcategory = "Partons"
    elif truthtype in ["WBosons", "ZBosons", "HBosons", "TQuarksFinal"]:
        truthcategory = "BosonTop"
        toolProperties['ParticleType'] = truthtype
    else:
        truthcategory = "FlavourLabel"
        toolProperties['ParticleType'] = truthtype
 
    tooltype = truthpartoptions[truthcategory]["ToolType"]
    toolProperties.update( PtMin = truthpartoptions[truthcategory]["ptmin"],
                           OutputName = "TruthLabel"+truthtype)
    ctp = tooltype("truthpartcopy_"+truthtype,
                   **toolProperties
                   )
    return ctp
 
# Generates input truth particle containers for truth jets
def getCopyTruthJetParticles(modspec, cflags):
    truthclassif = getMCTruthClassifier()
 
    truthpartcopy = CompFactory.CopyTruthJetParticles(
        "truthpartcopy"+modspec,
        OutputName="JetInputTruthParticles"+modspec,
        MCTruthClassifier=truthclassif)
    if modspec=="NoWZ":
        truthpartcopy.IncludePromptLeptons=False
        truthpartcopy.IncludePromptPhotons=False
        truthpartcopy.IncludeMuons=True
        truthpartcopy.IncludeNeutrinos=True
    if modspec=="DressedWZ":
        truthpartcopy.IncludePromptLeptons=False
        truthpartcopy.IncludePromptPhotons=True
        truthpartcopy.IncludeMuons=True
        truthpartcopy.IncludeNeutrinos=True
        truthpartcopy.DressingDecorationNames=['TruthParticles.dressedPhoton_e','TruthParticles.dressedPhoton_mu']
        
        truthpartcopy.ExtraInputs = {( 'xAOD::TruthParticleContainer' , 'StoreGateSvc+TruthParticles.dressedPhoton_e' ),
                                     ( 'xAOD::TruthParticleContainer' , 'StoreGateSvc+TruthParticles.dressedPhoton_mu' )}
    if modspec=="Charged":
        truthpartcopy.ChargedParticlesOnly=True
    return truthpartcopy
 
 
def getCopyTruthJetParticlesGEN(modspec, cflags):
    """  Build truth constituents as in EVTGEN jobs in the r21 config.
    IMPORTANT : this is expected to be temporary, only to reproduce the EVTGEN r21 config with the new config. The definitions should be harmonized with reco-level at some point and this function removed.
    The source for r21 EVTGEN config was in GeneratorFilters/share/common/GenerateTruthJets.py
    """
    truthclassif = getMCTruthClassifier()
 
    if modspec == "":
        return CompFactory.CopyTruthJetParticles("truthpartcopy",
                                                 OutputName="JetInputTruthParticlesGEN",
                                                 MCTruthClassifier=truthclassif)
    elif modspec=="NoWZ":
         return CompFactory.CopyTruthJetParticles("truthpartcopywz",
                                                  OutputName="JetInputTruthParticlesGENNoWZ",
                                                  MCTruthClassifier=truthclassif,
                                                  IncludePromptLeptons=False)
 
 
def _getCommonLabelNames(prefix):
    """Internal unlity to name labels
 
    Returns a dictionary to configure labeling tools. Takes one
    argument which is prefixed to each label.
    """
    return dict(
        LabelName=f"{prefix}TruthLabelID",
        DoubleLabelName=f"{prefix}ExtendedTruthLabelID",
        LabelPtName=f"{prefix}TruthLabelPt",
        LabelLxyName=f"{prefix}TruthLabelLxy",
        LabelDRName=f"{prefix}TruthLabelDR",
        LabelPdgIdName=f"{prefix}TruthLabelPdgId",
        LabelPositionDPhiName=f"{prefix}TruthLabelPositionDPhi",
        LabelPositionDEtaName=f"{prefix}TruthLabelPositionDEta",
        LabelBarcodeName=f"{prefix}TruthLabelBarcode",
        ChildLxyName=f"{prefix}TruthLabelChildLxy",
        ChildPtName=f"{prefix}TruthLabelChildPt",
        ChildPdgIdName=f"{prefix}TruthLabelChildPdgId",
        ChildPositionDPhiName=f"{prefix}TruthLabelChildPositionDPhi",
        ChildPositionDEtaName=f"{prefix}TruthLabelChildPositionDEta",
    )
 
 
# ATLASRECTS-8290: remove use_barcode here
def getJetDeltaRFlavorLabelTool(name='jetdrlabeler', jet_pt_min=5000, collection="Final", dr_max=0.3, use_barcode=False):
    """Get the standard flavor tagging delta-R labeling tool
 
    Uses cone matching to B, C and tau truth particles.
    """
    prefix_to_name = "HadronConeExcl"
    if collection != "Final":
        prefix_to_name += collection
        name+=collection
 
    return CompFactory.ParticleJetDeltaRLabelTool(
        name,
        **_getCommonLabelNames(prefix_to_name),
        BLabelName = "ConeExclBHadrons"+collection,
        CLabelName = "ConeExclCHadrons"+collection,
        TauLabelName = "ConeExclTausFinal",
        BParticleCollection = "TruthLabelBHadrons"+collection,
        CParticleCollection = "TruthLabelCHadrons"+collection,
        TauParticleCollection = "TruthLabelTausFinal",
        PartPtMin = 5000.,
        DRMax = dr_max,
        MatchMode = "MinDR",
        JetPtMin = jet_pt_min,
        useBarcode=use_barcode, # ATLASRECTS-8290: remove this eventually
        )
 
 
def getJetQuarkChargeLabelTool(name='jetquarkcharge',
                               output_label = "QuarkChargeTruthLabelID",
                               hadron_label = "HadronGhostInitialTruthLabelPdgId",
                               parton_label = "PartonExtendedTruthLabelID"):
    """Get the flavor tagging labeling tool to store the quark charge of a jet
    It is necessary to have saved before an hadron label and a parton label, otherwise the code will break
    """
 
    import ParticleJetTools.quarkChargeMap as qcMap
    charge_map = qcMap.hadrons_dict
    
    return CompFactory.JetQuarkChargeLabelingTool(
        name,
        HadronDecorationName = hadron_label,
        PartonDecorationName = parton_label,
        HadronChargeMap = charge_map,
        OutputName = output_label,
        )
 
def getJetDeltaRLabelTool(jetdef, modspec):
    """returns a ParticleJetDeltaRLabelTool
    Cone matching for B, C and tau truth for all but track jets.
 
    This function is meant to be used as callback from JetRecConfig where
    it is called as func(jetdef, modspec). Hence the jetdef argument even if not used in this case.
    """
    jetptmin = float(modspec)
    name = "jetdrlabeler_jetpt{0}GeV".format(int(jetptmin/1000))
    return getJetDeltaRFlavorLabelTool(name, jetptmin)
 
 
def getJetQuarkChargeTool(jetdef, modspec):
    """returns a JetQuarkChargeLabelingTool
 
    This function is meant to be used as callback from JetRecConfig where
    it is called as func(jetdef, modspec). Hence the jetdef argument even if not used.
    """
    if modspec == "":
        return getJetQuarkChargeLabelTool()
 
    outputLabel,hadronLabel,partonLabel = modspec.split(",")
    
    name = f"jetquarkcharge_{outputLabel}_{hadronLabel}_{partonLabel}"
    return getJetQuarkChargeLabelTool(name,output_label = outputLabel, hadron_label = hadronLabel, parton_label = partonLabel)
 
def getJetDeltaRInitialLabelTool(jetdef, modspec):
    """returns a ParticleJetDeltaRLabelTool
    Cone matching for B, C and tau truth for all but track jets.
 
    This function is meant to be used as callback from JetRecConfig where
    it is called as func(jetdef, modspec). Hence the jetdef argument even if not used in this case.
    """
    jetptmin = float(modspec)
    name = "jetdrlabeler_jetpt{0}GeV".format(int(jetptmin/1000))
    return getJetDeltaRFlavorLabelTool(name, jetptmin, collection = "Initial")
 
 
# ATLASRECTS-8290: remove use_barcode here
def getJetGhostFlavorLabelTool(
        name="jetghostlabeler",
        collection="Final",
        use_barcode=False):
 
    prefix_to_name = "HadronGhost"
    if collection != "Final":
        prefix_to_name += collection
        name+=collection
    return CompFactory.ParticleJetGhostLabelTool(
        name,
        **_getCommonLabelNames(prefix_to_name),
        GhostBName = "GhostBHadrons"+collection,
        GhostCName = "GhostCHadrons"+collection,
        GhostTauName = "GhostTausFinal",
        useBarcode = use_barcode, # ATLASRECTS-8290: remove this eventually
        PartPtMin = 5000.0
    )
 
def getJetGhostInitialLabelTool(jetdef, modspec):
    """get ghost-based flavor tagging labeling
 
    This is a wrapper for JetRecConfig where it's called as
    func(jetdef, modspec)
    """
    return getJetGhostFlavorLabelTool(modspec,collection = "Initial")
 
def getJetGhostLabelTool(jetdef, modspec):
    """get ghost-based flavor tagging labeling
 
    This is a wrapper for JetRecConfig where it's called as
    func(jetdef, modspec)
    """
    return getJetGhostFlavorLabelTool(modspec)
 
 
def getJetTruthLabelTool(jetdef, modspec):
 
    isTruthJet = 'Truth' in jetdef.fullname()
 
    if not isinstance(modspec, str):
        raise ValueError("JetTruthLabelingTool can only be scheduled with str as modspec")
    else:
        truthLabel = str(modspec)
 
    jetTruthLabelTool = CompFactory.JetTruthLabelingTool('truthlabeler_{0}'.format(truthLabel),
                                                         RecoJetContainer = jetdef.fullname(),
                                                         IsTruthJetCollection = isTruthJet,
                                                         TruthLabelName = truthLabel)
 
    return jetTruthLabelTool
 
def getJetTruthLabelToolPrereqs(jetdef, modspec):
    return ["input:AntiKt10TruthDressedWZSoftDropBeta100Zcut10Jets"] if modspec == "R10WZTruthLabel_R22v1" and jetdef._cflags.Input.isMC else []
 
def getJetPileupLabelTool(jetdef, modspec):
 
    jetPileupLabelTool = CompFactory.JetPileupLabelingTool('pileuplabeler',
                                                           RecoJetContainer = jetdef.fullname(),
                                                           TruthJetContainer= "AntiKt4TruthDressedWZJets")
 
    return jetPileupLabelTool