StandardJetMods.py

Go to the documentation of this file.

# Copyright (C) 2002-2025 CERN for the benefit of the ATLAS collaboration
"""
This module defines the standard JetModifier tools used in jet reco
 
Definitions are grouped in a dictionary of tool configurations using the helpers defined
in package configs.
This dict maps a modifier alias to the JetModifier config object
that in turn will be responsible for generating a configured tool.
 
The JetModifier config class is defined in JetDefinition.py 
 
 Args to the JetModifier constructor are:
   1. Tool Type (ignored if the helper is a custom one)
   2. Tool Name (ignored if the helper is a custom one)
   3. createfn : helper function which build the actual tool. If none, we just instantiate the tool type. 
   4. prereqs : Prerequisites  (default to []). Can also be a function which returns a list
   X. all other keyword arguments are directly interpreted as Property of the tool. 
       for ex, passing 'PtMin=10.' will configure the tool as in 'tool.PtMin = 10'
       we can pass function as the value : 
         'JetContainerName=nameFunc' will configure as in 'tool.JetContainerName=nameFunc(jetdef, modspec)'
 
        --> should this be by default? prefer to avoid ignored args
"""
from .JetDefinition import JetModifier
from .Utilities import ldict
from .StandardJetConstits import inputsFromContext
 
from JetRecConfig.JetRecCommon import isMC
 
stdJetModifiers = ldict()
 
 
# Many JetMoment tools need to know the name of the container they operate on.
# We set the function below as the 'JetContainer' property so the config system
# can assign the right name to the c++ tool.
def _jetname(jetdef,modspec):
    return jetdef.fullname()
 
 
def _constitContainername(jetdef,modspec):
    """Returns the exact constituent container name used to build jetdef"""
    constitdef = jetdef.inputdef
    containername = constitdef.containername(jetdef).split(':')[-1] if callable(constitdef.containername) else constitdef.containername
    return containername
 
 

stdJetModifiers.update( 
    Sort   = JetModifier("JetSorter","jetsort"),
    Filter = JetModifier("JetFilterTool","jetptfilter_{modspec}",
                         # we give a function as PtMin : it will be evaluated when instantiating the tool (modspec is specified with this tool
                         # alias like "Filter:10000" --> PtMin=100000).
                         PtMin = lambda jdef,modspec: int(modspec) 
                         ),
    Filter_ifnotESD = JetModifier("JetFilterTool","jetptfilter_{modspec}",
                                 PtMin = lambda _,modspec: 1,
                                 ),
    # Filter that can be easily turned off via pre-exec for e.g. PHYSVAL
    Filter_calibThreshold = JetModifier("JetFilterTool","jetptfilter_{modspec}",
                                        PtMin = lambda jetdef,modspec: 1 if not jetdef._cflags.Jet.useCalibJetThreshold else int(modspec),
                                       )
)
 

 
# Calibration
try:
    from JetCalibTools import JetCalibToolsConfig
    stdJetModifiers.update(
        Calib = JetModifier("JetCalibrationTool","jetcalib_jetcoll_calibseq",
                            createfn=JetCalibToolsConfig.getJetCalibToolFromString,
                            prereqs=lambda mod,jetdef : JetCalibToolsConfig.getJetCalibToolPrereqs(mod,jetdef)+[inputsFromContext("Vertices")])
    )
 
    from JetCalibTools import JetResponseToolConfig
    stdJetModifiers.update(
        # *******************
        # Response variables
        Response = JetModifier("JetResponseTool","response",
                             JetContainer=_jetname,
                             prereqs= ["input:EventDensity"],
                             createfn=JetResponseToolConfig.getJetResponseTool,
                             ),
    )
 
except ModuleNotFoundError:
    from AthenaCommon import Logging
    jetlog = Logging.logging.getLogger('JetStandardMods')
    jetlog.info("No JetMomentTools pakage found (expected in AthGeneration or analysis releases). Some jet calculations will be disabled.")
    pass
 
try:
    from JetCalibTools import JetCalibToolsCfg
    stdJetModifiers.update(
        CalibNew = JetModifier("JetCalibTool","jetcalib_jetcoll_calibseq",
                               createfn=JetCalibToolsCfg.defineJetCalibTool,
                               prereqs=lambda mod,jetdef : JetCalibToolsCfg.getJetCalibToolPrereqs(jetdef,mod)+[inputsFromContext("Vertices")])
    )
except ModuleNotFoundError:
    from AthenaCommon import Logging
    jetlog = Logging.logging.getLogger('JetStandardMods')
    jetlog.info("No JetMomentTools pakage found (expected in AthGeneration or analysis releases). Some jet calculations will be disabled.")
    pass
 
# TBD:
# All items below in principle will support decoration mode, rather
# than only non-const modification. Mode of operation should be
# determined by interface called from parent tool/alg.
 
 
# Many JetMoment tools need to know the name of the container they operate on.
# We set the function below as the 'JetContainer' property so the config system
# can assign the right name to the c++ tool.
def _jetname(jetdef,modspec):
    return jetdef.fullname()
 
 
# Standard jet moments
try:
    from JetMomentTools import JetMomentToolsConfig
    stdJetModifiers.update(
 
        # Easy cases, no special config or prereqs, just default tool config 
        ClusterMoments =  JetModifier("JetClusterMomentsTool", "clsmoms", JetContainer = _jetname),
        ECPSFrac =        JetModifier("JetECPSFractionTool", "ecpsfrac", JetContainer = _jetname),
        Width =           JetModifier("JetWidthTool", "width", JetContainer = _jetname),
 
        # More complex cases here
        CaloEnergies =    JetModifier("JetCaloEnergies", "jetens", 
                                      prereqs=["mod:EMScaleMom"], 
                                      Calculations=["EMFrac", "HECFrac", "PSFrac", "FracSamplingMax"], JetContainer = _jetname,
                                      ),
        
        CaloEnergiesLargeR =    JetModifier("JetCaloEnergies", "jetenslargeR", 
                                      prereqs=["mod:EMScaleMom"], 
                                      Calculations=["EMFrac", "HECFrac", "PSFrac", "EM3Frac", "Tile0Frac", "EffNClusts"], JetContainer = _jetname,
                                      calcClusterBasedVars = True,
                                      ),
 
        # CaloEnergiesClus is only relevant for FE-based jet collections
        CaloEnergiesClus = JetModifier("JetCaloEnergies", "jetensclus",
                                       prereqs=["mod:EMScaleMom"], 
                                       Calculations=["EMFrac", "HECFrac", "PSFrac", "FracSamplingMax"], JetContainer = _jetname,
                                       calcClusterBasedVars = True),
 
        CaloQuality =     JetModifier("JetCaloQualityTool", "caloqual",
                                      TimingCuts = [5,10],
                                      Calculations = ["LArQuality", "N90Constituents", "NegativeE", "Timing", "HECQuality", "Centroid", "AverageLArQF", "BchCorrCell"],JetContainer = _jetname),
 
        CaloQualityFE =   JetModifier("JetCaloQualityToolFE", "caloqualFE",
                                      TimingCuts = [5,10],
                                      ThresholdCuts = [90],
                                      Calculations = ["LArQuality", "NegativeE", "Timing", "HECQuality", "Centroid", "AverageLArQF", "BchCorrCell"],JetContainer = _jetname),
 
        N90 =             JetModifier("JetCaloQualityTool", "n90",
                                      Calculations = ["N90Constituents"],JetContainer = _jetname),
 
        ConstitFourMom =  JetModifier("JetConstitFourMomTool", "constitfourmom_basename",
                                      createfn=JetMomentToolsConfig.getConstitFourMomTool,),
        EMScaleMom =      JetModifier("JetEMScaleMomTool", "emscalemom_basename",
                                      createfn=JetMomentToolsConfig.getEMScaleMomTool,
                                      JetContainer = _jetname),
 
        JVF =             JetModifier("JetVertexFractionTool", "jvf",
                                      createfn=JetMomentToolsConfig.getJVFTool,
                                      prereqs = [inputsFromContext("TVA"), "mod:TrackMoments", inputsFromContext("Vertices")],
                                      JetContainer = _jetname),
        JVT =             JetModifier("JetVertexTaggerTool", "jvt",
                                       createfn=JetMomentToolsConfig.getJVTTool,
                                       prereqs = [ "mod:JVF" ],JetContainer = _jetname),
        NNJVT =           JetModifier("JetVertexNNTagger", "nnjvt",
                                       createfn=JetMomentToolsConfig.getNNJvtTool,
                                       prereqs = [ "mod:JVF" ],JetContainer = _jetname),
        LArHVCorr =       JetModifier("JetLArHVTool", "larhvcorr",
                                       prereqs = ["mod:EMScaleMom"],JetContainer = _jetname),
        OriginSetPV =     JetModifier("JetOriginCorrectionTool", "origin_setpv",
                                       prereqs = [ "mod:JVF" ],JetContainer = _jetname, OnlyAssignPV=True),
        TrackMoments =    JetModifier("JetTrackMomentsTool", "trkmoms",
                                      createfn=JetMomentToolsConfig.getTrackMomentsTool,
                                      prereqs = [ inputsFromContext("TVA"),"ghost:Track" ],JetContainer = _jetname),
 
        TrackSumMoments = JetModifier("JetTrackSumMomentsTool", "trksummoms",
                                      createfn=JetMomentToolsConfig.getTrackSumMomentsTool,
                                      prereqs = [ inputsFromContext("TVA"),"ghost:Track" ],JetContainer = _jetname),
        Charge =          JetModifier("JetChargeTool", "jetcharge", 
                                      prereqs = [ "ghost:Track" ]),
 
        QGTagging =       JetModifier("JetQGTaggerVariableTool", "qgtagging",
                                      createfn=JetMomentToolsConfig.getQGTaggingTool,
                                      prereqs = lambda _,jetdef :
                                      [inputsFromContext("TVA"),"mod:TrackMoments"] +
                                      (["mod:JetPtAssociation"] if not isMC(jetdef._cflags) else []),
                                      JetContainer = _jetname),
 
        fJVT =           JetModifier("JetForwardPFlowJvtTool", "fJVT",
                                     createfn=JetMomentToolsConfig.getPFlowfJVTTool,
                                     prereqs = [inputsFromContext("TVA"),inputsFromContext("EventDensity"),inputsFromContext("Vertices"),"mod:NNJVT"],
                                     JetContainer = _jetname),
 
        bJVT =           JetModifier("JetBalancePFlowJvtTool", "bJVT",
                                     createfn=JetMomentToolsConfig.getPFlowbJVTTool,
                                     prereqs = [inputsFromContext("EventDensity"),inputsFromContext("Vertices"),"mod:NNJVT"],
                                     JetContainer = _jetname),
 
        ConstitFrac =    JetModifier("JetConstituentFrac", "constitFrac",
                                     JetContainer = _jetname,
                                     ),
        
        groomMRatio =    JetModifier("JetGroomMRatio", "groomMRatio",
                                     JetContainer = _jetname,
                                     ),
 
        JetPtAssociation = JetModifier("JetPtAssociationTool", "jetPtAssociation",
                                       filterfn=isMC,
                                       createfn=JetMomentToolsConfig.getJetPtAssociationTool,
                                       prereqs=["ghost:Truth"],
                                       JetContainer = _jetname
                                       ),
 
        numConstit = JetModifier("JetNumConstitTool", "numConstit",
                                 JetContainer = _jetname,
                                 ),
 
        # *******************
        # Isolation variables
        jetiso = JetModifier("JetIsolationTool","iso",
                             JetContainer=_jetname,
                             InputConstitContainer = _constitContainername,
                             IsolationCalculations = ["IsoFixedCone:5:Pt",   "IsoFixedCone:5:PtPUsub",],
                             RhoKey = lambda jetdef, specs : "Kt4"+jetdef.inputdef.label+"EventShape" ,
                             prereqs= ["input:EventDensity"], #lambda spec,jetdef : ["input:Kt4"+jetdef.inputdef.label+"EventShape",],
                             ),
 
        jetisoTruth = JetModifier("JetIsolationTool","isoTruth",
                                   JetContainer=_jetname,
                                   InputConstitContainer = _constitContainername,
                                   IsolationCalculations = ["IsoFixedCone:5:Pt", ],
                                   RhoKey = "",
                                   ),
    )
 
except ModuleNotFoundError:
    from AthenaCommon import Logging
    jetlog = Logging.logging.getLogger('JetStandardMods')
    jetlog.info("No JetMomentTools pakage found (expected in AthGeneration or analysis releases). Some jet calculations will be disabled.")
    pass
 
 
# Truth labelling moments
from ParticleJetTools import ParticleJetToolsConfig
stdJetModifiers.update(
    # Easy cases, no special config or prereqs, just default tool config
    PartonTruthLabel = JetModifier("Analysis::JetPartonTruthLabel","partontruthlabel",
                                    prereqs=["ghost:Partons"]
                                   ),
 
    # More complex cases here
    JetDeltaRLabel =   JetModifier("ParticleJetDeltaRLabelTool","jetdrlabeler_jetptmin",
                                   createfn=ParticleJetToolsConfig.getJetDeltaRLabelTool,
                                   prereqs=["ghost:BHadronsFinal",
                                            "ghost:CHadronsFinal",
                                            "ghost:TausFinal"]
                                   ),
 
 
    JetGhostLabel =    JetModifier("ParticleJetGhostLabelTool","jetghostlabeler",
                                   createfn=ParticleJetToolsConfig.getJetGhostLabelTool,
                                   prereqs=["ghost:BHadronsFinal",
                                            "ghost:CHadronsFinal",
                                            "ghost:TausFinal"]
                                   ),
 
    
    JetDeltaRInitialLabel =   JetModifier("ParticleJetDeltaRLabelTool","jetdrlabelerinitial_jetptmin",
                                          createfn=ParticleJetToolsConfig.getJetDeltaRInitialLabelTool,
                                          prereqs=["ghost:BHadronsInitial",
                                                   "ghost:CHadronsInitial",
                                                   "ghost:TausFinal"]
                                   ),
 
    
    JetGhostInitialLabel =    JetModifier("ParticleJetGhostLabelTool","jetghostinitiallabeler",
                                          createfn=ParticleJetToolsConfig.getJetGhostInitialLabelTool,
                                          prereqs=["ghost:BHadronsInitial",
                                                   "ghost:CHadronsInitial",
                                                   "ghost:TausFinal"]
                                   ),
 
    JetQuarkChargeLabel =    JetModifier("JetQuarkChargeLabelingTool","jetquarkchargetool",
                                         createfn=ParticleJetToolsConfig.getJetQuarkChargeTool,
                                         prereqs=["mod:JetGhostInitialLabel","mod:JetGhostLabel","mod:PartonTruthLabel"]
                                         ),
 
 
    JetTaggingTruthLabel = JetModifier("JetTaggingTruthLabel", "truthlabeler_{mods}",
                                       filterfn=isMC,
                                       createfn=ParticleJetToolsConfig.getJetTruthLabelTool,
                                       prereqs=lambda modspec,jetdef: ParticleJetToolsConfig.getJetTruthLabelToolPrereqs(jetdef, modspec)
                                      ),
 
    JetPileupLabel = JetModifier("JetPileupLabel", "pileuplabeler_{mods}",
                                 filterfn=isMC,
                                 createfn=ParticleJetToolsConfig.getJetPileupLabelTool,
                                 prereqs=["input:AntiKt4TruthDressedWZJets"]
                                 )
    )
 
 
 
# Substructure tools 
stdJetModifiers.update( 
    nsubjettiness = JetModifier( "NSubjettinessTool", "nsubjettiness", Alpha = 1.0,
                                 JetContainer = _jetname),
    nsubjettinessR = JetModifier( "NSubjettinessRatiosTool", "nsubjettinessR",),
 
    
    ktdr       = JetModifier("KtDeltaRTool", "ktdr", JetRadius = 0.4),
 
    ktsplitter = JetModifier( "KTSplittingScaleTool", "ktsplitter", JetContainer = _jetname),
    
    angularity = JetModifier( "AngularityTool", "angularity"),
    
    dipolarity = JetModifier( "DipolarityTool", "dipolarity",SubJetRadius = 0.3),
    
    planarflow = JetModifier( "PlanarFlowTool", "planarflow"),
 
    ktmassdrop = JetModifier( "KtMassDropTool", "ktmassdrop"),
 
    ecorr      = JetModifier( "EnergyCorrelatorTool", "ecorr", Beta = 1.0, JetContainer = _jetname),
    ecorrR     = JetModifier( "EnergyCorrelatorRatiosTool", "ecorrR", ),
 
    ecorrgeneral = JetModifier( "EnergyCorrelatorGeneralizedTool", "ecorrgeneral", DoLSeries = True, JetContainer = _jetname),
    ecorrgeneralratios = JetModifier( "EnergyCorrelatorGeneralizedRatiosTool", "ecorrgeneralratios",  DoLSeries = True),
 
    comshapes = JetModifier( "CenterOfMassShapesTool", "comshapes", JetContainer = _jetname),
 
    pull      = JetModifier("JetPullTool", "pull",  UseEtaInsteadOfY = False, IncludeTensorMoments = True ),
 
    charge    = JetModifier( "JetChargeTool", "charge", K=1.0),
 
    qw = JetModifier( "QwTool", "qw", JetContainer = _jetname),
 
    softdropobs = JetModifier("SoftDropObservablesTool", "softdropobs"),
)
 
# Substructure tagger tools: q/g
try :
    from JetMomentTools import JetMomentToolsConfig
    stdJetModifiers.update( 
        qgtransformer = JetModifier("BoostedJetTaggerTool", "qgtransformer",
                                    createfn=JetMomentToolsConfig.getBoostedJetTaggerToolQG,
                                    JetContainer = _jetname,
                                    SuppressInputDependence = True),
    )
except ModuleNotFoundError:
    from AthenaCommon import Logging
    jetlog = Logging.logging.getLogger('JetStandardMods')
    jetlog.info("No JetMomentTools pakage found (expected in AthGeneration or analysis releases). Some jet calculations will be disabled.")
    pass
 
# Substructure tagger tools: top
try :
    from JetMomentTools import JetMomentToolsConfig
    stdJetModifiers.update( 
        toptransformer = JetModifier("BoostedJetTaggerTool", "toptransformer",
                                     createfn=JetMomentToolsConfig.getBoostedJetTaggerToolTop,
                                     JetContainer = _jetname,
                                     SuppressInputDependence = True),
    )
except ModuleNotFoundError:
    from AthenaCommon import Logging
    jetlog = Logging.logging.getLogger('JetStandardMods')
    jetlog.info("No JetMomentTools pakage found (expected in AthGeneration or analysis releases). Some jet calculations will be disabled.")
    pass
 
# Substructure tagger tools: w
try :
    from JetMomentTools import JetMomentToolsConfig
    stdJetModifiers.update( 
        wtransformer = JetModifier("BoostedJetTaggerTool", "wtransformer",
                            createfn=JetMomentToolsConfig.getBoostedJetTaggerToolW,
                            JetContainer = _jetname,
                            SuppressInputDependence = True),
    )
except ModuleNotFoundError:
    from AthenaCommon import Logging
    jetlog = Logging.logging.getLogger('JetStandardMods')
    jetlog.info("No JetMomentTools pakage found (expected in AthGeneration or analysis releases). Some jet calculations will be disabled.")
    pass
 
# Substructure tagger tools: w mass dec
try :
    from JetMomentTools import JetMomentToolsConfig
    stdJetModifiers.update( 
        wtransformer_massdec = JetModifier("BoostedJetTaggerTool", "wtransformer_massdec",
                            createfn=JetMomentToolsConfig.getBoostedJetTaggerToolWMassDec,
                            JetContainer = _jetname,
                            SuppressInputDependence = True),
    )
except ModuleNotFoundError:
    from AthenaCommon import Logging
    jetlog = Logging.logging.getLogger('JetStandardMods')
    jetlog.info("No JetMomentTools pakage found (expected in AthGeneration or analysis releases). Some jet calculations will be disabled.")
    pass
 
# VR track-jet decorations
stdJetModifiers.update(
    vr = JetModifier( "FlavorTagDiscriminants::VRJetOverlapDecoratorTool", "vr")
)