HbbCommon.py 32.8 KB
Newer Older
1
# Copyright (C) 2002-2020 CERN for the benefit of the ATLAS collaboration
2
3


Dan Guest's avatar
Dan Guest committed
4
5
6
# import Common Algs
from DerivationFrameworkJetEtMiss.JetCommon import DFJetAlgs

7
from DerivationFrameworkCore.DerivationFrameworkMaster import (
Zach Marshall's avatar
Zach Marshall committed
8
9
10
    DerivationFrameworkHasTruth as hasMCTruth)
from DerivationFrameworkCore.DerivationFrameworkMaster import (
    DerivationFrameworkIsDataOverlay)
Dan Guest's avatar
Dan Guest committed
11
12
13
14
# I wish we didn't need this
from BTagging.BTaggingConfiguration import getConfiguration
ConfInst=getConfiguration()

15
from GaudiKernel.Configurable import WARNING, VERBOSE
Dan Guest's avatar
Dan Guest committed
16

Dan Guest's avatar
Dan Guest committed
17
# Import star stuff (it was like that when I got here)
18
19
20
21
22
23
24
25
26
from DerivationFrameworkJetEtMiss.JetCommon import *
from DerivationFrameworkJetEtMiss.ExtendedJetCommon import *
from JetRec.JetRecConf import JetAlgorithm

#===================================================================
# ExKt/CoM Subjets
#===================================================================

# make exkt subjet finding tool
27
def buildExclusiveSubjets(ToolSvc, JetCollectionName, subjet_mode, nsubjet, doGhostAssoc, doTrackSubJet, ExGhostLabels=["GhostBHadronsFinal", "GhostCHadronsFinal", "GhostTrack"], min_subjet_pt_mev = 0):
Jie Yu's avatar
Jie Yu committed
28
    #
29
    # a full list of ExGhostLabels = ["GhostBHadronsFinal", "GhostBHadronsInitial", "GhostBQuarksFinal", "GhostCHadronsFinal", "GhostCHadronsInitial",
Jie Yu's avatar
Jie Yu committed
30
31
    # "GhostCQuarksFinal", "GhostHBosons", "GhostPartons", "GhostTQuarksFinal", "GhostTausFinal", "GhostTruth", "GhostTrack"]
    #
32

Jie Yu's avatar
Jie Yu committed
33
    from JetRec.JetRecStandard import jtm
34
35
36
37
38
39

    subjetlabel = "Ex%s%iSubJets" % (subjet_mode, nsubjet)
    doCoM = False
    algj = "Kt"
    if subjet_mode == "CoM" :
      doCoM = True
Jie Yu's avatar
Jie Yu committed
40
41
42
43
      #
      #supported algorithms: Reconstruction/Jet/JetSubStructureUtils/Root/SubjetFinder.cxx
      #algorithms:"ee_kt" or "EEKt" for EE_Kt
      #           "cambridge" or "CamKt" for Cambridge
44
      #           "FastJetPlugin" for EECambridge plugin
Jie Yu's avatar
Jie Yu committed
45
      algj = "ee_kt"
46

47
48
49
50
51
    talabel = ""
    if doGhostAssoc:
      talabel = "GA"
    subjetlabel = "Ex%s%i%sSubJets" % (subjet_mode, nsubjet, talabel)

52
    # removing truth labels if runining on data
Zach Marshall's avatar
Zach Marshall committed
53
54
    if not hasMCTruth:
        ExGhostLabels = ["GhostTrack"]
55

56
57
58
    SubjetContainerName = "%sEx%s%i%sSubJets" % (JetCollectionName.replace("Jets", ""), subjet_mode, nsubjet, talabel)
    ExKtbbTagToolName = str( "Ex%s%sbbTagTool%i_%s" % (subjet_mode, talabel, nsubjet, JetCollectionName) )

Jie Yu's avatar
Jie Yu committed
59
60
61
62
63
64
65
66
67
    if hasattr(jtm, ExKtbbTagToolName):
        ExKtbbTagToolInstance = jtm[ ExKtbbTagToolName ]
        print " ExKtbbTagTool ", ExKtbbTagToolName, "is alredy in jtm.tools"
    else:
        print " building ExKtbbTagTool ", ExKtbbTagToolName

        from JetSubStructureMomentTools.JetSubStructureMomentToolsConf import SubjetFinderTool
        from JetSubStructureMomentTools.JetSubStructureMomentToolsConf import SubjetRecorderTool

68
        subjetrecorder = SubjetRecorderTool("subjetrecorder_%s%i%s_%s" % (subjet_mode, nsubjet, talabel, JetCollectionName))
Jie Yu's avatar
Jie Yu committed
69
70
71
72
73
74
75
76
77
78
79
80
        ToolSvc += subjetrecorder
        subjetrecorder.SubjetLabel = subjetlabel
        subjetrecorder.SubjetContainerName = SubjetContainerName

        if not hasattr(ToolSvc, "ExKtTrackSelectorLoose"):
          from InDetTrackSelectionTool.InDetTrackSelectionToolConf import InDet__InDetTrackSelectionTool
          ToolSvc += InDet__InDetTrackSelectionTool("ExKtTrackSelectorLoose",
                                                    CutLevel = "Loose",
                                                   )

        from JetTagTools.JetTagToolsConf import Analysis__ExKtbbTagTool
        ExKtbbTagToolInstance = Analysis__ExKtbbTagTool(
81
          name = ExKtbbTagToolName,
Jie Yu's avatar
Jie Yu committed
82
83
          JetAlgorithm = algj,
          JetRadius = 10.0,
84
          PtMin = min_subjet_pt_mev,
Jie Yu's avatar
Jie Yu committed
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
          ExclusiveNJets = nsubjet,
          doTrack = doTrackSubJet,
          InputJetContainerName = JetCollectionName,
          SubjetContainerName = SubjetContainerName,
          SubjetRecorder = subjetrecorder,
          SubjetLabel = subjetlabel,
          SubjetAlgorithm_BTAG = "AntiKt",
          SubjetRadius_BTAG = 0.4,
          TrackSelectionTool = ToolSvc.ExKtTrackSelectorLoose,
          PrimaryVtxContainerName = "PrimaryVertices",
          SubjetBoostConstituent = doCoM,
          GhostLabels = ",".join(ExGhostLabels)
        )
        jtm.add( ExKtbbTagToolInstance )
        ToolSvc += ExKtbbTagToolInstance
100
101
102
103

    return (ExKtbbTagToolInstance, SubjetContainerName)

#===================================================================
104
# Build ExKt or CoM Subjets, default is ExKt
105
#===================================================================
106
def addExKtCoM(sequence, ToolSvc, JetCollectionExCoM, nSubjets, doTrackSubJet, doGhostAssoc=False, ExGhostLabels=["GhostBHadronsFinal","GhostCHadronsFinal","GhostTrack"], min_subjet_pt_mev = 0, subjetAlgName = "Kt"):
mkhader's avatar
mkhader committed
107
108
109
    if(subjetAlgName != "Kt" and subjetAlgName != "CoM"):
      print "WARNING:  Subjet type must be Kt or CoM.  Using Kt as default!"
      subjetAlgName = "Kt"
Jie Yu's avatar
Jie Yu committed
110
    from JetRec.JetRecStandard import jtm
111
    ExCoMJetCollection__SubJet = []
112

113
    (ExCoMbbTagToolInstance, SubjetContainerName) = buildExclusiveSubjets(ToolSvc, JetCollectionExCoM, subjetAlgName, nSubjets, doGhostAssoc, doTrackSubJet, ExGhostLabels, min_subjet_pt_mev)
114

115
    ExCoMJetCollection__SubJet += [SubjetContainerName]
116

Jie Yu's avatar
add CoM    
Jie Yu committed
117

118
119
120
121
122
    excomELresetName = "ELreset_subjet_%s" % (SubjetContainerName.replace("Jets", ""))
    excomELresetNameLJet = "ELreset_Large_%sjet_%s" % (SubjetContainerName.replace("Jets", ""),JetCollectionExCoM.replace("Jets", ""))
    excomAlgName = "jfind_%s" % (SubjetContainerName)
    excomJetRecToolName = "%s" % (SubjetContainerName)
    excomBTagName = "BTagging_%s" % (SubjetContainerName.replace("Jets", ""))
Jie Yu's avatar
add CoM    
Jie Yu committed
123

124
125
    if excomAlgName in DFJetAlgs:
        print " Algorithm Ex%s "%subjetAlgName, excomAlgName, "already built before sequence ", sequence
Jie Yu's avatar
Jie Yu committed
126

127
128
        if hasattr(sequence, excomAlgName):
            print " sequence Ex%s "%subjetAlgName, sequence, "already has an instance of algorithm", excomAlgName
Jie Yu's avatar
Jie Yu committed
129
        else:
130
131
            print " Add algorithm Ex%s "%subjetAlgName, excomAlgName, "to sequence", sequence
            sequence += DFJetAlgs[excomAlgName]
Jie Yu's avatar
Jie Yu committed
132
            sequence += CfgMgr.xAODMaker__ElementLinkResetAlg(excomELresetName, SGKeys=[SubjetContainerName+"Aux."])
133
            sequence += DFJetAlgs[excomAlgName+"_btag"]
Jie Yu's avatar
Jie Yu committed
134
            sequence += CfgMgr.xAODMaker__ElementLinkResetAlg(excomELresetNameLJet, SGKeys=[JetCollectionExCoM+"Aux."])
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
    else:
        print " Algorithm Ex%s "%subjetAlgName, excomAlgName, " to be built sequence ", sequence
        if hasattr(jtm, excomJetRecToolName):
            print " Ex%sJetRecTool "%subjetAlgName, excomJetRecToolName, " already built sequence ",  sequence
            jetrec = jtm [ excomJetRecToolName ]
        else:
            print " Ex%s tool "%subjetAlgName, excomJetRecToolName, " to be built sequence ", sequence
            from JetRec.JetRecConf import JetRecTool
            jetrec = JetRecTool(
                                 name = excomJetRecToolName,
                                 OutputContainer = JetCollectionExCoM,
                                 InputContainer = JetCollectionExCoM,
                                 JetModifiers = [ExCoMbbTagToolInstance],
                               )
            jtm.add( jetrec )
            ToolSvc += jetrec

        excomJetRecBTagToolName = str( "%s_sub"%excomJetRecToolName )
        if hasattr(jtm, excomJetRecBTagToolName):
            print " Ex%sJetRecBTagTool "%subjetAlgName, excomJetRecBTagToolName, " already built sequence ",  sequence
            jetrec_btagging = jtm [ excomJetRecBTagToolName ]
        else:
            print " Ex%sJetRecBTagTool "%subjetAlgName, excomJetRecBTagToolName, " to be built sequence ",  sequence

            #make the btagging tool for excom jets
            from BTagging.BTaggingFlags import BTaggingFlags
            btag_excom = ConfInst.setupJetBTaggerTool(ToolSvc, JetCollection=excomJetRecToolName.replace("Jets", ""), AddToToolSvc=True, Verbose=True,
                         options={"name"         : excomBTagName.lower(),
                                  "BTagName"     : excomBTagName,
                                  "BTagJFVtxName": "JFVtx",
                                  "BTagSVName"   : "SecVtx",
                                  },
                         SetupScheme = "",
                         TaggerList = BTaggingFlags.StandardTaggers,
169
                         TrackAssociatorName="GhostTrack" if doGhostAssoc else "MatchedTracks"
170
171
172
173
174
                         )
            # running association + b-tagging on subjets now
            from BTagging.BTaggingConfiguration import comTrackAssoc, comMuonAssoc, defaultTrackAssoc, defaultMuonAssoc
            mods = [defaultTrackAssoc, defaultMuonAssoc, btag_excom]
            if(subjetAlgName=="CoM"): mods = [comTrackAssoc, comMuonAssoc, btag_excom]
Zach Marshall's avatar
Zach Marshall committed
175
            if hasMCTruth:
Jie Yu's avatar
Jie Yu committed
176
                mods.append(jtm.jetdrlabeler)
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202

            jetrec_btagging = JetRecTool( name = excomJetRecBTagToolName,
                                      InputContainer  = SubjetContainerName,
                                      OutputContainer = SubjetContainerName,
                                      JetModifiers    = mods)
            jtm.add( jetrec_btagging )
            ToolSvc += jetrec_btagging

        jetalg_excom = JetAlgorithm(
                                 name = excomAlgName,
                                 Tools = [jetrec],
                                )
        sequence += jetalg_excom
        DFJetAlgs[excomAlgName] = jetalg_excom

        # Reset EL for ExCoM subjets after all of them are built
        # Otherwise crashing for CoM during TrackToJetAssociation due to failure of finding the large-R parent jet.
        sequence += CfgMgr.xAODMaker__ElementLinkResetAlg(excomELresetName, SGKeys=[SubjetContainerName+"Aux."])

        jetalg_excom_btag = JetAlgorithm(
                                 name = excomAlgName+"_btag",
                                 Tools = [jetrec_btagging],
                                )
        sequence += jetalg_excom_btag
        DFJetAlgs[excomAlgName+"_btag"] = jetalg_excom_btag
        sequence += CfgMgr.xAODMaker__ElementLinkResetAlg(excomELresetNameLJet, SGKeys=[JetCollectionExCoM+"Aux."])
203

Jie Yu's avatar
Jie Yu committed
204
    return ExCoMJetCollection__SubJet
Jie Yu's avatar
add CoM    
Jie Yu committed
205

206
##################################################################
207
208
209
210
211
212
213
214
215
# Associate the VR track jets to existing large-R jet collections
# 
# For any groomed jet collections, the association will be done to
# the parent jet collection (so for the
# AntiKt10LCTopoTrimmedPtFrac5SmallR20Jets collection the 
# association will be done to the AntiKt10LCTopoJets collection).
# 
# If no jet collection is supplied, the association is done to
# AntiKt10LCTopoJets
216
##################################################################
217
def addVRJets(sequence, largeRColls = None, do_ghost=False, logger=None, doFlipTagger=False, training='201810', *pos_opts, **opts):
Dan Guest's avatar
Dan Guest committed
218
219
    from AthenaCommon import Logging

220
    if logger is None:
Dan Guest's avatar
Dan Guest committed
221
222
223
224
225
226
227
        logger = Logging.logging.getLogger('VRLogger')

    # define constants here, since we don't want every derivaiton
    # deciding for themselves what a VR jet is. If we do want this
    # flexibility, this code will need some rewriting to ensure that
    # there are no issues with train safety.
    if opts or pos_opts:
228
229
230
231
232
233
        logger.warning('Options specified for VR jets, they will be ignored')

    if largeRColls is None:
      largeRColls = [
          "AntiKt10LCTopoTrimmedPtFrac5SmallR20Jets"
          ]
Dan Guest's avatar
Dan Guest committed
234

235
    VRName, ghostLab = buildVRJets(sequence, do_ghost, logger, doFlipTagger, training)
236
237
238
239
240
241
242
243
244
    toAssociate = {
        ghostLab : ghostLab.lower()
        }
    
    ungroomedNames = []
    for collection in largeRColls:
      ungroomedName, getters = linkVRJetsToLargeRJets(sequence, collection, toAssociate)
      ungroomedNames.append(ungroomedName)

Ines Ochoa's avatar
Ines Ochoa committed
245
246
247
248
    return dict(
        ungroomedNames=ungroomedNames,
        VRJetName=VRName,
        ghostJetLinkName=ghostLab)
249

250
251
252
253
254
255
def buildVRJets(sequence, do_ghost, logger = None, doFlipTagger=False, training='201810'):

    from AthenaCommon import Logging

    if logger is None:
        logger = Logging.logging.getLogger('VRLogger')
256

257
    supported_trainings = ['201810', '201903']
258
259
260
    # Check allowed trainings
    # Is there a better way to do this with a central DB?
    if training not in ['201810', '201903']:
261
262
      logger.warning("WARNING: Using an unsupported training tag! This is UNDEFINED and will probably break. Please choose a training tag from")
      logger.warning(supported_trainings)
263

264
265
    from JetRec.JetRecStandard import jtm

266
267
    # Making Chris Happy: all VR track-jet b-tagging should have the training campaign label
    trainingTag = '_BTagging%s' % (training)
268
269
270

    VRJetName="AntiKtVR30Rmax4Rmin02Track%s" % (trainingTag)
    VRGhostLabel="GhostVR30Rmax4Rmin02TrackJet%s" % (trainingTag)
Dan Guest's avatar
Dan Guest committed
271
272
273
274
    VRJetAlg="AntiKt"
    VRJetRadius=0.4
    VRJetInputs='pv0track'
    VRJetOptions = dict(
275
        ghostArea = 0 , ptmin = 4000,
Dan Guest's avatar
Dan Guest committed
276
277
278
        variableRMinRadius = 0.02, variableRMassScale = 30000,
        calibOpt = "none")

279
280
281
282
283
284
    # Change some options if we have do_ghost set to true. Hopefully
    # this will be the only VR collection in the future.
    if do_ghost:
        ghost_suffix = "GhostTag"
        VRJetName += ghost_suffix
        VRGhostLabel += ghost_suffix
285
286
287
288
289

    #==========================================================
    # Build VR jets
    #==========================================================

290
291
    from DerivationFrameworkJetEtMiss.ExtendedJetCommon import nameJetsFromAlg
    VRJetRecToolName = nameJetsFromAlg(VRJetName)
292
293
294
    VRJetAlgName = "jfind_%s" % (VRJetRecToolName)
    VRJetBTagName = "BTagging_%s" % (VRJetName.replace('BTagging',''))

295
296
    logger.info("VR Btag name: %s" % VRJetBTagName)
    logger.info("VR jet name: %s" % VRJetRecToolName)
297

298
299
    from AthenaCommon.AppMgr import ToolSvc

300
    #make the btagging tool for VR jets
Chris Pollard's avatar
Chris Pollard committed
301
    from BTagging.BTaggingFlags import BTaggingFlags
302
    BTaggingFlags.CalibrationChannelAliases += ["AntiKtVR30Rmax4Rmin02Track->AntiKtVR30Rmax4Rmin02Track,AntiKt4EMTopo"]
303
    BTaggingFlags.CalibrationChannelAliases += ["%s->AntiKtVR30Rmax4Rmin02Track,AntiKt4EMTopo" % (VRJetName)]
Dan Guest's avatar
Dan Guest committed
304
305
306
307
308
309
310
311
    btag_vrjets = ConfInst.setupJetBTaggerTool(
        ToolSvc, JetCollection=VRJetRecToolName, AddToToolSvc=True, Verbose=True,
        options={"name"         : VRJetBTagName.lower(),
                 "BTagName"     : VRJetBTagName,
                 "BTagJFVtxName": "JFVtx",
                 "BTagSVName"   : "SecVtx",
        },
        SetupScheme = "",
Binbin Dong's avatar
Binbin Dong committed
312
        TaggerList = BTaggingFlags.ExpertTaggers if doFlipTagger else BTaggingFlags.StandardTaggers,
313
        TrackAssociatorName="GhostTrack" if do_ghost else "MatchedTracks",
Dan Guest's avatar
Dan Guest committed
314
    )
315

Dan Guest's avatar
Dan Guest committed
316
317
318
319
320
321
    # add delta-R to nearest jet
    from FlavorTagDiscriminants.FlavorTagDiscriminantsLibConf import (
        FlavorTagDiscriminants__VRJetOverlapDecoratorTool as VRORTool )
    vrdr_label = VRORTool(name=VRJetRecToolName + "_VRLabeling")
    ToolSvc += vrdr_label

322
323
324
325
326
327
328
    # add Ghost label id
    from ParticleJetTools.ParticleJetToolsConf import (
        ParticleJetGhostLabelTool as GhostLabelTool)
    gl_tool = GhostLabelTool(
        name=VRJetRecToolName + "_GhostLabeling")
    ToolSvc += gl_tool

329
330
    from BTagging.BTaggingConfiguration import defaultTrackAssoc, defaultMuonAssoc

331
332
333
    # Slice the array - this forces a copy so that if we modify it we don't also
    # change the array in jtm.
    pseudoJetGetters = jtm.gettersMap[VRJetInputs][:]
334
335
336
337

    # We want to include ghost associated tracks in the pv0 tracks so that
    # we can use the looser ghost association criteria for b-tagging.
    if VRJetInputs == "pv0track":
338
        pseudoJetGetters.append(jtm["gtrackget"])
339

340
    if VRJetAlgName in DFJetAlgs:
341
        logger.info("Algorithm %s already built before" % VRJetAlgName)
342
343

        if hasattr(sequence, VRJetAlgName):
344
            logger.info("Sequence %s already has an instance of algorithm %s" % (sequence, VRJetAlgName))
345
        else:
346
            logger.info("Add algorithm %s to sequence %s" % (VRJetAlgName, sequence))
347
348
            sequence += DFJetAlgs[VRJetAlgName]
    else:
349
        logger.info("Create algorithm %s" % VRJetAlgName)
350
351

        if hasattr(jtm, VRJetRecToolName):
352
            logger.info("JetRecTool %s is alredy in jtm.tools in sequence %s" % (VRJetRecToolName, sequence))
353
        else:
354
            logger.info("Create JetRecTool %s" % VRJetRecToolName)
355
            #can only run trackjetdrlabeler with truth labels, so MC only
356

Dan Guest's avatar
Dan Guest committed
357
            mods = [defaultTrackAssoc, defaultMuonAssoc, btag_vrjets,
358
                    vrdr_label]
359

Zach Marshall's avatar
Zach Marshall committed
360
            if hasMCTruth:
361
                mods += [jtm.trackjetdrlabeler, gl_tool]
362

363
364
365
366
367
368
            jtm.addJetFinder(
                VRJetRecToolName,
                VRJetAlg,
                VRJetRadius,
                pseudoJetGetters,
                modifiersin=mods,
Dan Guest's avatar
Dan Guest committed
369
                ivtxin=0,
370
                **VRJetOptions)
371
372
373
374
375
376
377
378
379
380
381

        from JetRec.JetRecConf import JetAlgorithm
        jetalg_smallvr30_track = JetAlgorithm(VRJetAlgName, Tools = [ jtm[VRJetRecToolName] ])
        sequence += jetalg_smallvr30_track
        DFJetAlgs[VRJetAlgName] = jetalg_smallvr30_track

    #==========================================================
    # Build PseudoJet Getter
    #==========================================================

    pjgettername = VRGhostLabel.lower()
382
    from DerivationFrameworkJetEtMiss.ExtendedJetCommon import nameJetsFromAlg
383
384

    if hasattr(jtm, pjgettername):
385
        logger.info("Found %s in jtm in sequence %s" % (pjgettername, sequence))
386
    else:
387
        logger.info("Add %s to jtm in sequence %s" % (pjgettername, sequence))
388
389
        
        inputContainerName = jetFlags.containerNamePrefix() + nameJetsFromAlg(VRJetName)
390

391

392
393
394
        from JetRec.JetRecConf import PseudoJetGetter
        jtm += PseudoJetGetter(
          pjgettername,                                                          # give a unique name
395
          InputContainer = inputContainerName,                                   # SG key
396
397
398
399
400
          Label = VRGhostLabel,                                                  # this is the name you'll use to retrieve ghost associated VR track jets
          OutputContainer = "PseudoJet" + VRGhostLabel,
          SkipNegativeEnergy = True,
          GhostScale = 1.e-20,                                                   # this makes the PseudoJet Ghosts, and thus the reco flow will treat them as such
        )
401
    return VRJetName, VRGhostLabel
402

403
##################################################################
Xuanhong Lou's avatar
Xuanhong Lou committed
404
405
406
# Build variable-R calorimeter jets
##################################################################
def addVRCaloJets(sequence,outputlist,dotruth=True,writeUngroomed=False):
Zach Marshall's avatar
Zach Marshall committed
407
    if hasMCTruth and dotruth:
Xuanhong Lou's avatar
Xuanhong Lou committed
408
409
410
411
412
413
414
        addTrimmedJets('AntiKt', 1.0, 'Truth', rclus=0.2, ptfrac=0.05, variableRMassScale=600000, variableRMinRadius=0.2, mods="truth_groomed",
                       algseq=sequence, outputGroup=outputlist, writeUngroomed=writeUngroomed)
    addTrimmedJets('AntiKt', 1.0, 'PV0Track', rclus=0.2, ptfrac=0.05, variableRMassScale=600000, variableRMinRadius=0.2, mods="groomed",
                   algseq=sequence, outputGroup=outputlist, writeUngroomed=writeUngroomed)
    addTrimmedJets('AntiKt', 1.0, 'LCTopo', rclus=0.2, ptfrac=0.05, variableRMassScale=600000, variableRMinRadius=0.2, mods="lctopo_groomed",
                   algseq=sequence, outputGroup=outputlist, writeUngroomed=writeUngroomed)

415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
###################################################################
## Utils: get jetRectoTool
###################################################################
def getJetRecTool(collection, getParent=True):
  """Get the JetRecTool for a given collection from the jtm. If getParent is
     set then if that tool has an InputContainer property set then try and 
     retrieve the JetRecTool for that parent, continue going until there is no
     InputContainer property. Will raise a KeyError if no JetRecTool can be
     found at any stage.
  """
  from JetRec.JetRecStandardToolManager import jtm
  try:
    jetRecTool = jtm[collection]
  except KeyError as e:
     raise KeyError("JetRecTool {0} not present in jtm".format(collection) )
  if getParent and hasattr(jetRecTool, "InputContainer") and jetRecTool.InputContainer:
    jetRecTool = getJetRecTool(jetRecTool.InputContainer, True)
  return jetRecTool

###################################################################
## Utils: link jets (copied over from Jon Burr in EXOT27Utils)
###################################################################
437
def linkVRJetsToLargeRJets(
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
    sequence, collection, getters):
  """Re-run jet finding for a jet collection using a new list of 
    PseudoJetGetters. These PseudoJetGetters should already have been loaded
    into jtm.
    collection should be the name of the jet collection, which should already
    have been sequenced, so it's jet rec tool will exist in jtm.
    getters should be a map of PseudoJetGetters, each key being the name of the
    (non ghost) collection, with the value being the name of the 
    PseudoJetGetter in jtm.
    Returns the name of the ungroomed collection that is the parent of
    'collection' (this will be the same as 'collection' if this isn't groomed)
    *and* the list of ghost labels (these are the element link names).
  """
  from JetRec.JetRecStandardToolManager import jtm
  import DerivationFrameworkJetEtMiss.JetCommon as JetCommon
453
  from DerivationFrameworkJetEtMiss.ExtendedJetCommon import nameJetsFromAlg
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
  logger = Logging.logging.getLogger('HbbTaggerLog')
  # First, retrieve the original JetRecTool - this is the one that made the
  # *ungroomed* jets, not the groomed ones. Ghost association is done to
  # ungroomed objects
  originalJetRecTool = getJetRecTool(collection, getParent = True)
  originalUngroomedName = originalJetRecTool.name()
  ungroomedJetAlg = originalUngroomedName
  if ungroomedJetAlg.endswith("Jets"):
    ungroomedJetAlg = ungroomedJetAlg[:-4]
  originalFinder = jtm[originalJetRecTool.JetFinder.getName()]
  originalGetters = [jtm[g.getName()] for g in originalJetRecTool.PseudoJetGetters]
  newGetters = [jtm[g] for g in getters.values()]

  # Next, prepare the names of the new objects we'll need from jtm

  comb_name = "_".join(getters.keys() )
  LargeRJetFindingAlg = "jfind_{0}_{1}".format(collection, comb_name).lower()
  LargeRJetPrefix     = "{0}_{1}".format(collection, comb_name)
472
  LargeRJets = nameJetsFromAlg(LargeRJetPrefix)
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
  LinkTransferAlg     = "LinkTransfer_{0}_{1}".format(collection, comb_name)

  # Check to see if this large R jet collection is already known to JetCommon
  if LargeRJetFindingAlg in JetCommon.DFJetAlgs:
    logger.info("Found {0} in DFJetAlgs".format(LargeRJetFindingAlg) )
    # Is it in our sequence?
    if hasattr(sequence, LargeRJetFindingAlg):
      logger.info("Algorithm already exists in the input sequence. Will not "
          "add again")
    else:
      logger.info("Adding algorithm into the sequence {0}".format(sequence) )
      sequence += JetCommon.DFJetAlgs[LargeRJetFindingAlg]
  else:
    # Check to see if the corresponding JetRecTool already exists
    if hasattr(jtm, LargeRJets):
      logger.info("JetRecTool {0} already exists in jtm".format(LargeRJets) )
    else:
      logger.info("Create a new JetRecTool {0}".format(LargeRJets) )
      JetCommon.OutputJets.setdefault("CustomJets", []).append(LargeRJets)
      originalModifiers = [jtm[m.getName()] for m in originalJetRecTool.JetModifiers]
      jtm.addJetFinder(
          output = LargeRJets,
          alg = originalFinder.JetAlgorithm,
          radius = originalFinder.JetRadius,
          gettersin = originalGetters + newGetters,
          modifiersin = originalModifiers,
          ghostArea = 0,
          ptmin = originalFinder.PtMin,
          variableRMinRadius = originalFinder.VariableRMinRadius,
          variableRMassScale = originalFinder.VariableRMassScale,
          calibOpt = "none")
      # Note that we don't need ptminFilter as this was included in the original
      # list of JetModifiers
    logger.info(
        "Creating new jet algorithm {0} and adding it to sequence {1}".format(
          LargeRJetFindingAlg, sequence) )
    theJetAlg = JetAlgorithm(LargeRJetFindingAlg, Tools = [jtm[LargeRJets] ])
    sequence += theJetAlg
    JetCommon.DFJetAlgs[LargeRJetFindingAlg] = theJetAlg

  # Everything so far has been to create the links on a copy of the ungroomed
  # collection. Now we need to copy those links over to the original ungroomed
  # collection.
  from DerivationFrameworkJetEtMiss.ExtendedJetCommon import getJetExternalAssocTool, applyJetAugmentation
  assocTool = getJetExternalAssocTool(
      ungroomedJetAlg,
      LargeRJetPrefix,
      MomentPrefix = '',
      ListOfOldLinkNames=[g.Label for g in newGetters]
      )
  applyJetAugmentation(
      ungroomedJetAlg, LinkTransferAlg, sequence, assocTool)

  return originalUngroomedName, [g.Label for g in newGetters]



530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
#===================================================================
# Utils: Copy Jets
#===================================================================
def addCopyJet(FTAG5Seq, ToolSvc, InputJetCollectionName, OutputJetCollectionName, **options):
  doShallowCopy = options.pop("ShallowCopy", False)

  from JetRec.JetRecConf import JetRecTool
  jetrec = JetRecTool(
                       name = "JetRecTool_CopyJet_%s_%s" % (InputJetCollectionName, OutputJetCollectionName),
                       OutputContainer = OutputJetCollectionName,
                       InputContainer = InputJetCollectionName,

                       ShallowCopy = doShallowCopy,
                       **options
                     )

  ToolSvc += jetrec
  FTAG5Seq += JetAlgorithm(
                            name = "JetAlgorithm_CopyJet_%s_%s" % (InputJetCollectionName, OutputJetCollectionName),
                            Tools = [jetrec],
                          )

  return OutputJetCollectionName
Dan Guest's avatar
Dan Guest committed
553
554
555
556
557


#========================================================================
# Hbb Tagger
#========================================================================
558
559
560
561
562
563
564
565
566
567
568
569
570
def get_unique_name(strings):
    clean_strings = []
    for string in strings:
        chars = []
        for char in string:
            chars += char if (char.isalpha() or char.isdigit()) else '_'
        clean_strings.append(''.join(chars))
    return '_'.join(clean_strings)

def addHbbTagger(
        sequence, ToolSvc, logger=None,
        output_level=WARNING,
        jet_collection="AntiKt10LCTopoTrimmedPtFrac5SmallR20",
Dan Guest's avatar
Dan Guest committed
571
572
        nn_file_name="BoostedJetTaggers/HbbTagger/Summer2018/Apr13HbbNetwork.json",
        nn_config_file="BoostedJetTaggers/HbbTaggerDNN/PreliminaryConfigNovember2017.json"):
573
574
575
    if logger is None:
        logger = Logging.logging.getLogger('HbbTaggerLog')

576
    fat_calibrator_name = get_unique_name(["HbbCalibrator", jet_collection])
Zach Marshall's avatar
Zach Marshall committed
577
578
579
580
    is_data = not hasMCTruth
    if DerivationFrameworkIsDataOverlay:
        raise RuntimeError('Do not know how to treat overlay data, cannot run the Hbb tagger')

Dan Guest's avatar
Dan Guest committed
581
582
583
    if not hasattr(ToolSvc, fat_calibrator_name):
        fatCalib = CfgMgr.JetCalibrationTool(
            fat_calibrator_name,
584
            OutputLevel=output_level,
Dan Guest's avatar
Dan Guest committed
585
586
587
588
589
590
591
592
593
            JetCollection=jet_collection,
            ConfigFile="JES_MC16recommendation_FatJet_JMS_comb_19Jan2018.config",
            CalibSequence="EtaJES_JMS",
            CalibArea="00-04-81",
            IsData=is_data)
        ToolSvc += fatCalib
        logger.info('set up {}'.format(fatCalib))
    else:
        logger.info('took {} from tool svc'.format(fat_calibrator_name))
Dan Guest's avatar
Dan Guest committed
594
        fatCalib = getattr(ToolSvc, fat_calibrator_name)
Dan Guest's avatar
Dan Guest committed
595

596
597
    # short name for naming tools
    nn_short_file = nn_file_name.split('/')[-1].split('.')[0]
Dan Guest's avatar
Dan Guest committed
598

599
    hbb_tagger_name = get_unique_name(["HbbTagger",nn_short_file])
Dan Guest's avatar
Dan Guest committed
600
601
602
    if not hasattr(ToolSvc, hbb_tagger_name):
        hbbTagger = CfgMgr.HbbTaggerDNN(
            hbb_tagger_name,
603
            OutputLevel=output_level,
Dan Guest's avatar
Dan Guest committed
604
605
            neuralNetworkFile=nn_file_name,
            configurationFile=nn_config_file)
Dan Guest's avatar
Dan Guest committed
606
607
608
609
        ToolSvc += hbbTagger
        logger.info('set up {}'.format(hbbTagger))
    else:
        logger.info('took {} from tool svc'.format(hbb_tagger_name))
Dan Guest's avatar
Dan Guest committed
610
        hbbTagger = getattr(ToolSvc, hbb_tagger_name)
Dan Guest's avatar
Dan Guest committed
611

612
613
    tagger_alg_name = get_unique_name(
        ["HbbTaggerAlg",jet_collection, nn_short_file])
Dan Guest's avatar
Dan Guest committed
614
615
616
617
618
619
620
    if not hasattr(sequence, tagger_alg_name):
        if tagger_alg_name in DFJetAlgs:
            sequence += DFJetAlgs[tagger_alg_name]
            logger.info('took {} from jet algs'.format(tagger_alg_name))
        else:
            tagger_alg = CfgMgr.HbbTaggingAlgorithm(
                tagger_alg_name,
621
                OutputLevel=output_level,
Dan Guest's avatar
Dan Guest committed
622
623
624
625
626
627
628
629
630
631
632
633
                jetCollectionName=(jet_collection + "Jets"),
                minPt=250e3,
                maxEta=2.0,
                tagger=hbbTagger,
                calibrationTool=fatCalib)
            DFJetAlgs[tagger_alg_name] = tagger_alg
            sequence += tagger_alg
            logger.info('set up {}'.format(tagger_alg))
    else:
        logger.info('{} already scheduled for {}'.format(
            tagger_alg_name, jet_collection))

634
def addRecommendedXbbTaggers(sequence, ToolSvc, logger=None):
635
636
637
638
639
640
    if logger is None:
        logger = Logging.logging.getLogger('addRecXbbLog')
    logger.warning('addRecommendedXbbTaggers decorates a deprecated XbbScore variable. This decoration will be skipped.')
    return

xbbTaggerExtraVariables = []
641

642
#====================================================================
643
# Large-R RC jets w/ ExKt 2 & 3 subjets
644
#===================================================================
645
def addExKtDoubleTaggerRCJets(sequence, ToolSvc):
646
647
   jetToolName = "DFReclustertingTool"
   algoName = "DFJetReclusteringAlgo"
648

649
   ExKtJetCollection__FatJet = "AntiKt8EMPFlowJets"
Yuan-Tang Chou's avatar
Yuan-Tang Chou committed
650
   ExKtJetCollection__SubJet = []
651

652
   if jetToolName not in DFJetAlgs:
653
     ToolSvc += CfgMgr.JetReclusteringTool(jetToolName,InputJetContainer="AntiKt4EMPFlowJets", OutputJetContainer="AntiKt8EMPFlowJets")
654
     getattr(ToolSvc,jetToolName).ReclusterRadius = 0.8
655
     getattr(ToolSvc,jetToolName).InputJetPtMin = 15
656
     getattr(ToolSvc,jetToolName).RCJetPtMin = 1
657
     getattr(ToolSvc,jetToolName).TrimPtFrac = 0
658
659
660
     getattr(ToolSvc,jetToolName).DoArea = False
     getattr(ToolSvc,jetToolName).GhostTracksInputContainer = "InDetTrackParticles"
     getattr(ToolSvc,jetToolName).GhostTracksVertexAssociationName  = "JetTrackVtxAssoc"
Zach Marshall's avatar
Zach Marshall committed
661
     if hasMCTruth:
662
663
       getattr(ToolSvc,jetToolName).GhostTruthBHadronsInputContainer = "BHadronsFinal"
       getattr(ToolSvc,jetToolName).GhostTruthCHadronsInputContainer = "CHadronsFinal"
664

665
666
     sequence += CfgMgr.AthJetReclusteringAlgo(algoName, JetReclusteringTool = getattr(ToolSvc,jetToolName))
     DFJetAlgs[jetToolName] = getattr(ToolSvc,jetToolName)
667

668
   # build subjets
669
   GhostLabels = ["GhostTrack"]
Zach Marshall's avatar
Zach Marshall committed
670
   if hasMCTruth:
671
672
     GhostLabels += ["GhostBHadronsFinal"]
     GhostLabels += ["GhostCHadronsFinal"]
673
   # N=2 subjets
674
   ExKtJetCollection__SubJet += addExKtCoM(sequence, ToolSvc, ExKtJetCollection__FatJet, nSubjets=2, doTrackSubJet=True, ExGhostLabels=GhostLabels, min_subjet_pt_mev=1)
675
   # N=3 subjets
676
   ExKtJetCollection__SubJet += addExKtCoM(sequence, ToolSvc, ExKtJetCollection__FatJet, nSubjets=3, doTrackSubJet=True, ExGhostLabels=GhostLabels, min_subjet_pt_mev=1)
677
678
679
680
681
   #Ghosttracks
   # N=2 subjets
   ExKtJetCollection__SubJet += addExKtCoM(sequence, ToolSvc, ExKtJetCollection__FatJet, nSubjets=2, doTrackSubJet=True, doGhostAssoc=True, ExGhostLabels=GhostLabels, min_subjet_pt_mev=1)
   # N=3 subjets
   ExKtJetCollection__SubJet += addExKtCoM(sequence, ToolSvc, ExKtJetCollection__FatJet, nSubjets=3, doTrackSubJet=True, doGhostAssoc=True, ExGhostLabels=GhostLabels, min_subjet_pt_mev=1)
682

683
   sequence += CfgMgr.xAODMaker__ElementLinkResetAlg("ELReset_AfterSubjetBuild", SGKeys=[name+"Aux." for name in ExKtJetCollection__SubJet])
684

685
   from BTagging.BTaggingFlags import BTaggingFlags
686
687
   BTaggingFlags.CalibrationChannelAliases += [ jetname[:-4].replace("PV0", "")+"->AntiKt4EMTopo" for jetname in ExKtJetCollection__FatJet ]
   BTaggingFlags.CalibrationChannelAliases += [ jetname[:-4].replace("PV0", "")+"->AntiKt4EMTopo" for jetname in ExKtJetCollection__SubJet ]
688

689
   sequence += CfgMgr.xAODMaker__ElementLinkResetAlg("ELReset_AfterBtag", SGKeys=[name+"Aux." for name in ExKtJetCollection__SubJet])
690

691
def addExKtDoubleTaggerScore(sequence, ToolSvc,
692
                             sec_vtx_collection = 'SoftBVrtClusterTool_MSVTight_Vertices',
693
694
695
696
697
698
699
                             nn_file_name = '/cvmfs/atlas.cern.ch/repo/sw/database/GroupData/dev/BTagging/DeepsetXbbTagger/202011/nn-config.json',
                             nn_config_file = 'BoostedJetTaggers/DeepsetXbbTagger/test_config.json',
                             logger=None):

    dexter_tagger_name = 'DeepsetXbbTagger' 
    if logger is None:
        logger = Logging.logging.getLogger('DexterLog')
700
701

    if not hasattr(ToolSvc, dexter_tagger_name):
702
        dexterTagger = CfgMgr.DeepsetXbbTagger(
703
704
705
706
707
            dexter_tagger_name,
            KerasConfigFile = nn_file_name,
            ConfigFile = nn_config_file,
            secvtxCollection = sec_vtx_collection)
        ToolSvc += dexterTagger
708
709
710
    else:
        logger.info('took {} from tool svc'.format(dexter_tagger_name))
        dexterTagger = getattr(ToolSvc, dexter_tagger_name)