JetPresel.py

Go to the documentation of this file.

# Copyright (C) 2002-2025 CERN for the benefit of the ATLAS collaboration
 
from AthenaCommon.Logging import logging
log = logging.getLogger(__name__)
 
import re
 
from .JetRecoCommon import getJetCalibDefaultString, jetChainParts, etaRangeAbbrev, jetRecoDictToString
from ..Menu.SignatureDicts import JetChainParts_Default
from TriggerMenuMT.HLT.Config.ControlFlow.HLTCFTools import NoHypoToolCreated
 
 
from TrigHLTJetHypo.TrigJetHypoToolConfig import trigJetHypoToolFromDict
 
# Extract preselection reco dictionary
# This is 1:1 based on the main (tracking/PF) jet collection
def getPreselRecoDict(reco,roiftf=False):
 
    # Define a fixed preselection dictionary for prototyping -- we may expand the options
    preselRecoDict = {
        'recoAlg':reco,
        'constitType':'tc',
        'clusterCalib':'em',
        'constitMod':'',
        'trkopt':'notrk',
        'ionopt':'noion',
    }
    ''' #Here you can set custom calibrations for large-R preselections. If you set to LCW you'll get an issue though, as the trigger expects the *same* topocluster collection to be used in the preselection and in the PFlow stage with tracking. Therefore this would need to be adapted, but it might not be so easy...
        
    if preselRecoDict['recoAlg']=='a10': #Setting LC calibrations for large-R jets
        preselRecoDict['clusterCalib']='lcw'
    '''
    preselRecoDict.update({
        'jetCalib':getJetCalibDefaultString(preselRecoDict['recoAlg'],preselRecoDict['constitType'],preselRecoDict['clusterCalib'])
        if preselRecoDict['recoAlg']=='a4' else 'nojcalib'
    }) #Adding default calibration for corresponding chain
    
    # Overwriting tracking option to roiftf tracking
    if roiftf: preselRecoDict['trkopt'] = 'roiftf'
 
    preselRecoDict['jetDefStr'] = jetRecoDictToString(preselRecoDict)
 
    return preselRecoDict
 
# Find the preselection definition in the chainParts
# Simply taken from the last chainPart
def extractPreselection(fullChainDict):
    jparts = jetChainParts(fullChainDict['chainParts'])
    return jparts[-1]['trkpresel']
 
# Calo jet preselection hypo tool generator
# Translate the preselection expression in the main jet chainDict into a temporary chainDict
# that is only seen by the standard hypo tool generator, and used to return a configured
# hypo tool for the preselection step
def caloPreselJetHypoToolFromDict(flags, mainChainDict):
    return _preselJetHypoToolFromDict(flags, mainChainDict)
    
 
def roiPreselJetHypoToolFromDict(flags, mainChainDict):
    return _preselJetHypoToolFromDict(flags, mainChainDict,doTaggingSel=True)
 
 
def _preselJetHypoToolFromDict(flags, mainChainDict, doTaggingSel=False):
 
    preselChainDict = dict(mainChainDict)
    preselChainDict['chainParts']=[]
    trkpresel = extractPreselection(mainChainDict)
 
    # Get from the last chainPart in order to avoid to specify preselection for every leg
    #TODO: add protection for cases where the preselection is not specified in the last chainPart
    presel_matched = re.match(r'presel(?P<cut>(VETOMULT)?\d?\d?(Z[\d\D]+)?[jacf]\d?(HT)?[\d\D]+)', trkpresel)
    assert presel_matched is not None, "Impossible to match preselection pattern for self.trkpresel=\'{0}\'.".format(trkpresel)
    presel_cut_str = presel_matched.groupdict()['cut'] #This is the cut string you want to parse. For example 'presel2j50XXj40'
    
    usingDIPZ = bool(re.match(r'.*Z', presel_cut_str)) # Need to determine if there's DIPZ leg anywhere to enforce central jets across all calopresel legs
    if usingDIPZ:
        findSel = re.finditer(r'(?P<nJet>\d?[jacf])(?P<ptcut>\d+)', presel_cut_str) 
        findAllJets=[]
        findAllPts=[]
        for match in findSel:
            nJ = match.group("nJet")
            findAllJets.append('1'+nJ if len(nJ)==1 else nJ)
            findAllPts.append(int(match.group("ptcut")))
        nAllJets = sum(int(i[:-1]) for i in findAllJets)
        nCentralJets = sum(int(i[:-1]) if 'c' in i else 0 for i in findAllJets)        
        ptCut = min(int(i) for i in findAllPts)
        assert nAllJets == nCentralJets, "Your preselection has a DIPZ part but not only central jets were required. This isn't currently supported. Please investigate."
 
    preselCommonJetParts = dict(JetChainParts_Default)
    
    for ip,p in enumerate(presel_cut_str.split('XX')):
        hascalSel= bool(re.match(r'.*emf\w?\d+', p))
        # This appears to be very much a hack... we should just have separate
        # functions for with/without bjet or tau selections. -- Chris Pollard
        if not doTaggingSel:  # Removing b-jet and tau parts if b-jet presel is not requested
            p = re.sub(r'(b|bg|bgtwo|gntau|uht1tau)\d\d', '', p)
 
        hasBjetSel = bool(re.match(r'.*(b|bg|bgtwo)\d\d', p))
        hasTauSel = bool(re.match(r'.*(gntau|uht1tau)\d\d', p))
        hasDIPZsel = bool(re.match(r'.*Z', p))
 
        if hasDIPZsel and not doTaggingSel: continue # Skipping calopresel step when DIPZ is run
        if usingDIPZ and not hasDIPZsel and not hasBjetSel and doTaggingSel: continue # Skipping roiftf step only when running the calo selection leg (and if in the DIPZ scenario)
 
        assert not ( (hasBjetSel or hasDIPZsel) and not doTaggingSel), "Your jet preselection has a b-jet or DIPZ part but a calo-only preselection was requested instead. This should not be possible. Please investigate."        
 
        bmatches = r'(?P<btagger>(b|bg|bgtwo))(?P<bwp>\d\d)' if hasBjetSel else ""
        taumatches = r'(?P<tauid>(gntau|uht1tau))(?P<tauwp>\d\d)' if hasTauSel else ""
 
        pattern_to_test = r'(?P<veto>(VETOMULT))?' # veto flag #TODO veto pt as well
        pattern_to_test += r'(?P<mult>\d*)?(?P<region>[jacf])' # jet multiplicity and region and ptlow
        pattern_to_test += r'(?P<scenario>(HT))?(?P<cut>\d+)?' # scenario string # could be made more general
        pattern_to_test += bmatches
        pattern_to_test += taumatches
        pattern_to_test += r'emf(?P<emfc>\d+)' if hascalSel else ''
        if hasDIPZsel: pattern_to_test = r'(?P<scenario>Z)(?P<dipzwp>\d+)?(?P<prefilt>(MAXMULT\d+[jacf]?)?)'
        matched = re.match(pattern_to_test, p)
        assert matched is not None, "Impossible to extract preselection cut for \'{0}\' substring. Please investigate.".format(p)
        cut_dict = matched.groupdict()
        if hasDIPZsel: cut_dict['region'] = 'c'
        if hasDIPZsel: cut_dict['cut'] = ptCut
 
 
        # any missing keys k below need to have cut_dict[k] set to "".
        testkeys = \
            [ "mult" , "emfc" , "cut" , "cut_add" , "prefilt"
            , "bwp" , "btagger" , "dipzwp"
            , "tauwp"
            ]
 
 
        for k in testkeys:
            cut_dict.setdefault(k, "")
 
        veto = 'veto' in cut_dict
        if veto:
            veto = cut_dict['veto'] is not None
        mult = cut_dict['mult']
        region = cut_dict['region']
        scenario = cut_dict['scenario']
        cut = cut_dict['cut']
        dipzwp = cut_dict['dipzwp']
        emfc = cut_dict['emfc']
        prefilters = []
    
        if mult=='': mult='1'
        etarange = etaRangeAbbrev[region]
        if veto:
            scenario = cut_dict['veto'][4:]
            assert (not (scenario in ['HT','Z'])), "Veto preselection not yet supported for HT or Z scenarios. Please investigate."
        if scenario == "HT":
            hyposcenario=f'HT{cut}XX{etarange}'
            threshold='0'
            chainPartName=f'j0_{hyposcenario}'
        elif scenario == "Z":
            hyposcenario=f'Z{dipzwp}XX{nCentralJets}c{cut}'
            prefilt = cut_dict['prefilt']   
            if prefilt != '': prefilters.append(prefilt)
            threshold='0'
            chainPartName=f'j0_{hyposcenario}'
        elif scenario == "MULT":
            hyposcenario=f'MULT0mult{mult}XX{cut}ptXX{etarange}'
            threshold='0'
            chainPartName=f'j0_{hyposcenario}'
            log.info(f'Generated chainPartName {chainPartName}')
        else:
            hyposcenario='simple'
            threshold=cut
            chainPartName=f'{mult}j{cut}_{etarange}'
        
        if hasBjetSel:
            btagger = cut_dict['btagger']
            bwp = cut_dict['bwp']
 
            if btagger == 'bg':
                btagger = 'gnone'
            elif btagger =='bgtwo':
                btagger = 'bgntwo'
            elif btagger == 'b':
                btagger = 'bdips'
 
 
        tmpChainDict = dict(preselCommonJetParts)
 
        if hasTauSel:
            tauid = cut_dict["tauid"]
            tauwp = cut_dict["tauwp"]
 
            tmpChainDict[tauid] = tauwp+tauid
 
        tmpChainDict.update(
            {'L1threshold': 'FSNOSEED',
            'chainPartName': chainPartName,
            'multiplicity': mult,
            'threshold': threshold,
            'etaRange':etarange,
            'jvt':'',
            'clrsel': emfc,
            'bsel': f'{bwp}{btagger}' if hasBjetSel else "",
            'chainPartIndex': ip,
            'hypoScenario': hyposcenario,
            'prefilters': prefilters,
            }
        )
 
        preselChainDict['chainParts'] += [tmpChainDict] 
 
 
    # We need to pad by the legs not in the preselection expression
    # otherwise the ComboHypo does not find the corresponding
    # legs in the DecisionObject and kills the event
    jetlegs = sum([p['signature'] in ["Jet","Bjet"] for p in mainChainDict['chainParts']])
    padding = jetlegs-len(preselChainDict['chainParts'])
    if padding>0:
        preselChainDict['chainParts'][-1]['chainPartName']+='_SHARED'
        preselChainDict['chainParts'][-1]['tboundary']='SHARED'
        dummyLegPart = dict(preselCommonJetParts)
        dummyLegPart.update(
            {'L1threshold': 'FSNOSEED',
                'chainPartName': 'j0_SHARED',
                'multiplicity': '1',
                'threshold': '0',
                'jvt':'',
                'tboundary': 'SHARED',
                }
        )
        preselChainDict['chainParts'] += [dict(dummyLegPart) for i in range(padding)]
        # Last one is not permitted to be shared as there is nothing following
        preselChainDict['chainParts'][-1]['chainPartName']='j0'
        preselChainDict['chainParts'][-1]['tboundary']=''
    
    # Update certain keys to be consistent with the main dict
    # Also set the chainPart indices correctly from the main dict
    # These should be index-parallel now, as we only receive jet chainParts
    for porig,ppresel in zip(mainChainDict['chainParts'],preselChainDict['chainParts']):
        for key in ['chainPartIndex','signature']:
            ppresel[key] = porig[key]
    
    assert(len(preselChainDict['chainParts'])==len(mainChainDict['chainParts']))
    try:
        return trigJetHypoToolFromDict(flags, preselChainDict)
    except NoHypoToolCreated as nohypo:
        raise nohypo # We only generate the hypo tool for the first jet leg
    except Exception as e:
        log.error("Failure with preselection for chain %s",mainChainDict['chainName'])
        raise e