From ffeead94c4a9d963132a74684a64e44a136dcdfe Mon Sep 17 00:00:00 2001
From: Zach Marshall <ZLMarshall@lbl.gov>
Date: Wed, 11 Nov 2020 16:14:08 +0100
Subject: [PATCH] Enabling truth jets in TRUTH derivations

This gets truth jet building working again in truth derivations. The
state after this MR is:

-) The correct large-R jets are built and have the required decorations
-) The correst small-R jets are build, but they currently lack the
desired truth flavor decorations
-) As discussed elsewhere, some extra branches exist in truth DAODs made
with master, because branch dropping seems to not be working quite right
yet
-) There are errors coming from CutFlowSvc to do with a mis-match in the
number of weights. I'm not quite sure how to solve this one; it looks
like CutFlowSvc has moved and been re-worked in master, so I might ask
for a little help with this one.

Note that I've disabled most of the jet-related items for derivations
running off of AODs. That's because the truth jet building code
currently conflicts with what the jet group is running, and it's not yet
clear to me how to fix this all. I hope the jet group will clean up a
bit in the next few days/weeks, and then we'll be able to harmonize.
---
 .../python/MCTruthCommon.py                   | 62 +++++++++++--------
 .../python/TruthDerivationTools.py            |  6 +-
 .../src/TruthQGDecorationTool.cxx             | 11 ++--
 3 files changed, 45 insertions(+), 34 deletions(-)

diff --git a/PhysicsAnalysis/DerivationFramework/DerivationFrameworkMCTruth/python/MCTruthCommon.py b/PhysicsAnalysis/DerivationFramework/DerivationFrameworkMCTruth/python/MCTruthCommon.py
index 767aecbce65..2158b5d1261 100644
--- a/PhysicsAnalysis/DerivationFramework/DerivationFrameworkMCTruth/python/MCTruthCommon.py
+++ b/PhysicsAnalysis/DerivationFramework/DerivationFrameworkMCTruth/python/MCTruthCommon.py
@@ -59,6 +59,8 @@ def simplePJGetter(Label, InputContainer):
 
 # Helper for adding truth jet collections
 def addTruthJets(kernel=None, decorationDressing=None):
+    if not dfInputIsEVNT:
+        return
     # Ensure that we are adding it to something, and that we haven't run it already
     if kernel is None:
         from DerivationFrameworkCore.DerivationFrameworkMaster import DerivationFrameworkJob
@@ -74,33 +76,34 @@ def addTruthJets(kernel=None, decorationDressing=None):
 
         # Set up the copy truth jet particle algorithms
         from ParticleJetTools.ParticleJetToolsConf import CopyTruthJetParticles
-        ToolSvc += CopyTruthJetParticles("truthpartcopy",
+        ToolSvc += CopyTruthJetParticles("TruthPartCopy",
                                          OutputName="JetInputTruthParticles",
                                          MCTruthClassifier=ToolSvc.DFCommonTruthClassifier,BarCodeFromMetadata=barCodeFromMetadata)
-        ToolSvc += CopyTruthJetParticles("truthpartcopywz",
+        ToolSvc += CopyTruthJetParticles("TruthPartCopyWZ",
                                          OutputName="JetInputTruthParticlesNoWZ",
                                          MCTruthClassifier=ToolSvc.DFCommonTruthClassifier,BarCodeFromMetadata=barCodeFromMetadata,
                                          IncludePromptLeptons=False)
-        ToolSvc += CopyTruthJetParticles("truthpartcopydressedwz",
+        ToolSvc += CopyTruthJetParticles("TruthPartCopyDressedWZ",
                                          OutputName="JetInputTruthParticlesDressedWZ",
                                          MCTruthClassifier=ToolSvc.DFCommonTruthClassifier,
                                          IncludePromptLeptons=False,IncludePromptPhotons=False,
                                          IncludeMuons=True,IncludeNeutrinos=True,BarCodeFromMetadata=barCodeFromMetadata,
                                          FSRPhotonCone=-1., DressingDecorationName=decorationDressing)
-        ToolSvc += CopyTruthJetParticles("truthpartcopycharged", OutputName="JetInputTruthParticlesCharged",
+        ToolSvc += CopyTruthJetParticles("TruthPartCopyCharged", OutputName="JetInputTruthParticlesCharged",
                                          MCTruthClassifier=ToolSvc.DFCommonTruthClassifier,
                                          ChargedParticlesOnly=True,
                                          BarCodeFromMetadata=barCodeFromMetadata)
         from JetRec import JetRecConf
         kernel += JetRecConf.JetAlgorithm("MCTruthCommonJetTruthCopyAlg",
-                                          Tools=[ToolSvc.truthpartcopy,ToolSvc.truthpartcopywz,ToolSvc.truthpartcopydressedwz,ToolSvc.truthpartcopycharged])
+                                          Tools=[ToolSvc.TruthPartCopy,ToolSvc.TruthPartCopyWZ,
+                                                 ToolSvc.TruthPartCopyDressedWZ,ToolSvc.TruthPartCopyCharged])
 
         # Set up pseudo-jet getters
         from JetRec import JetRecConf
-        kernel += simplePJGetter( Label = "Truth", InputContainer = ToolSvc.truthpartcopy.OutputName )
-        kernel += simplePJGetter( Label = "TruthWZ", InputContainer = ToolSvc.truthpartcopywz.OutputName )
-        kernel += simplePJGetter( Label = "TruthDressedWZ", InputContainer = ToolSvc.truthpartcopydressedwz.OutputName )
-        kernel += simplePJGetter( Label = "TruthCharged", InputContainer = ToolSvc.truthpartcopycharged.OutputName )
+        kernel += simplePJGetter( Label = "Truth", InputContainer = ToolSvc.TruthPartCopy.OutputName )
+        kernel += simplePJGetter( Label = "TruthWZ", InputContainer = ToolSvc.TruthPartCopyWZ.OutputName )
+        kernel += simplePJGetter( Label = "TruthDressedWZ", InputContainer = ToolSvc.TruthPartCopyDressedWZ.OutputName )
+        kernel += simplePJGetter( Label = "TruthCharged", InputContainer = ToolSvc.TruthPartCopyCharged.OutputName )
 
         # Set up the jet builder (no area moments)
         from AthenaCommon import CfgMgr
@@ -204,11 +207,15 @@ def addTruthJets(kernel=None, decorationDressing=None):
                                                                             GhostArea = 0.01,
                                                                             PtMin = threshold
                                                                             )
+            from JetSubStructureMomentTools.JetSubStructureMomentToolsConf import EnergyCorrelatorTool
+            DFCommon_EnCorr = EnergyCorrelatorTool("DFCommon_EnCorr", Beta = 1.0)
+            from JetSubStructureMomentTools.JetSubStructureMomentToolsConf import NSubjettinessTool
+            DFCommon_NSubjettiness = NSubjettinessTool("DFCommon_NSubjettiness",Alpha = 1.0)
             AntiKt10TruthSoftDropBeta100Zcut10JetsRec = CfgMgr.JetRecTool("AntiKt10TruthSoftDropBeta100Zcut10JetsRec",
                                                                           JetGroomer = groomer,
                                                                           InputPseudoJets = [kernel.TruthGet.OutputContainer],
                                                                           OutputContainer = "AntiKt10TruthSoftDropBeta100Zcut10Jets",
-                                                                          JetModifiers = [ToolSvc.partontruthlabel],
+                                                                          JetModifiers = [ToolSvc.partontruthlabel,DFCommon_EnCorr,DFCommon_NSubjettiness],
                                                                           JetFinder = AntiKt10TruthSoftDropBeta100Zcut10JetsFinder)
             kernel += CfgMgr.JetAlgorithm("AntiKt10TruthSoftDropBeta100Zcut10JetsAlg",Tools=[AntiKt10TruthSoftDropBeta100Zcut10JetsRec])
 
@@ -282,20 +289,15 @@ def schedulePostJetMCTruthAugmentations(kernel=None, decorationDressing=None):
     augmentationToolsList = [ dfTruth.DFCommonTruthTauDressingTool ]
 
     #Save the post-shower HT and MET filter values that will make combining filtered samples easier (adds to the EventInfo)
-    #from DerivationFrameworkMCTruth.GenFilterToolSetup import DFCommonTruthGenFilter
-
-    # schedule the special truth building tools and add them to a common augmentation; note taus are handled separately below
-    #from DerivationFrameworkMCTruth.TruthDerivationTools import DFCommonTruthQGLabelTool
-    #augmentationToolsList += [ DFCommonTruthGenFilter,
-    #                          DFCommonTruthQGLabelTool]
-    augmentationToolsList = []
-    #if decorationDressing is not None:
-    #    from DerivationFrameworkMCTruth.DerivationFrameworkMCTruthConf import DerivationFramework__TruthQGDecorationTool
-    #    DFCommonTruthDressedWZQGLabelTool = DerivationFramework__TruthQGDecorationTool(name="DFCommonTruthDressedWZQGLabelTool",
-    #                                                              JetCollection = "AntiKt4TruthDressedWZJets")
-    #    from AthenaCommon.AppMgr import ToolSvc
-    #    ToolSvc += DFCommonTruthDressedWZQGLabelTool
-    #    augmentationToolsList += [ DFCommonTruthDressedWZQGLabelTool ]
+    if dfInputIsEVNT:
+        from DerivationFrameworkMCTruth.GenFilterToolSetup import DFCommonTruthGenFilter
+
+        # schedule the special truth building tools and add them to a common augmentation; note taus are handled separately below
+        from DerivationFrameworkMCTruth.TruthDerivationTools import DFCommonTruthDressedWZQGLabelTool
+        augmentationToolsList += [ DFCommonTruthGenFilter ]
+                                   #DFCommonTruthDressedWZQGLabelTool] - missing decoration from FTAG
+    else:
+        augmentationToolsList = []
     # SUSY signal decorations
     from DerivationFrameworkSUSY.DecorateSUSYProcess import IsSUSYSignal
     if IsSUSYSignal():
@@ -323,7 +325,7 @@ def addStandardTruthContents(kernel=None,
         from AthenaCommon.AppMgr import ToolSvc
         ToolSvc.DFCommonTruthTauDressingTool.decorationName=decorationDressing
     # Jets and MET
-    #addTruthJets(kernel, decorationDressing)
+    addTruthJets(kernel, decorationDressing)
     addTruthMET(kernel)
     # Tools that must come after jets
     schedulePostJetMCTruthAugmentations(kernel, decorationDressing)
@@ -331,7 +333,8 @@ def addStandardTruthContents(kernel=None,
     addTruthCollectionNavigationDecorations(kernel, ["TruthElectrons", "TruthMuons", "TruthPhotons", "TruthTaus", "TruthNeutrinos", "TruthBSM", "TruthBottom", "TruthTop", "TruthBoson"], prefix=prefix)
     # Some more additions for standard TRUTH3
     addBosonsAndDownstreamParticles(kernel)
-    #addLargeRJetD2(kernel)
+    if dfInputIsEVNT:
+        addLargeRJetD2(kernel)
     # Special collection for BSM particles
     addBSMAndDownstreamParticles(kernel)
     # Special collection for Born leptons
@@ -339,7 +342,8 @@ def addStandardTruthContents(kernel=None,
     # Special collection for hard scatter (matrix element) - save TWO extra generations of particles
     addHardScatterCollection(kernel,2)
     # Energy density for isolation corrections
-    #addTruthEnergyDensity(kernel)
+    if dfInputIsEVNT:
+        addTruthEnergyDensity(kernel)
 
 
 def addParentAndDownstreamParticles(kernel=None,
@@ -580,6 +584,8 @@ def addTruthEnergyDensity(kernel=None):
                                                             AbsRapidityMax      = 1.5,
                                                             OutputContainer     = "TruthIsoCentralEventShape",
                                                            )
+        # Note the helper function mangles the naming in a specific way that is not sufficiently general
+        DFCommonTruthCentralEDTool.InputContainer = kernel.TruthGet.OutputContainer
         ToolSvc += DFCommonTruthCentralEDTool
         kernel += EventDensityAthAlg("DFCommonTruthCentralEDAlg", EventDensityTool = DFCommonTruthCentralEDTool )
     if not hasattr(ToolSvc,'EDTruthForwardTool'):
@@ -590,6 +596,8 @@ def addTruthEnergyDensity(kernel=None):
                                                             AbsRapidityMax      = 3.0,
                                                             OutputContainer     = "TruthIsoForwardEventShape",
                                                            )
+        # Note the helper function mangles the naming in a specific way that is not sufficiently general
+        DFCommonTruthForwardEDTool.InputContainer = kernel.TruthGet.OutputContainer
         ToolSvc += DFCommonTruthForwardEDTool
         kernel += EventDensityAthAlg("DFCommonTruthForwardEDAlg", EventDensityTool = DFCommonTruthForwardEDTool )
 
diff --git a/PhysicsAnalysis/DerivationFramework/DerivationFrameworkMCTruth/python/TruthDerivationTools.py b/PhysicsAnalysis/DerivationFramework/DerivationFrameworkMCTruth/python/TruthDerivationTools.py
index 85e43033239..e7dc98771fd 100644
--- a/PhysicsAnalysis/DerivationFramework/DerivationFrameworkMCTruth/python/TruthDerivationTools.py
+++ b/PhysicsAnalysis/DerivationFramework/DerivationFrameworkMCTruth/python/TruthDerivationTools.py
@@ -198,6 +198,6 @@ DFCommonTruthPhotonIsolationTool3 = DerivationFramework__TruthIsolationTool(name
 ToolSvc += DFCommonTruthPhotonIsolationTool3
 # Quark/gluon decoration for jets
 from DerivationFrameworkMCTruth.DerivationFrameworkMCTruthConf import DerivationFramework__TruthQGDecorationTool
-DFCommonTruthQGLabelTool = DerivationFramework__TruthQGDecorationTool(name="DFCommonTruthQGLabelTool",
-                                                                  JetCollection = "AntiKt4TruthDressedWZJets")
-ToolSvc += DFCommonTruthQGLabelTool
+DFCommonTruthDressedWZQGLabelTool = DerivationFramework__TruthQGDecorationTool(name="DFCommonTruthDressedWZQGLabelTool",
+                                                                              JetCollection = "AntiKt4TruthDressedWZJets")
+ToolSvc += DFCommonTruthDressedWZQGLabelTool
diff --git a/PhysicsAnalysis/DerivationFramework/DerivationFrameworkMCTruth/src/TruthQGDecorationTool.cxx b/PhysicsAnalysis/DerivationFramework/DerivationFrameworkMCTruth/src/TruthQGDecorationTool.cxx
index 42540612ff5..eb47543ac15 100644
--- a/PhysicsAnalysis/DerivationFramework/DerivationFrameworkMCTruth/src/TruthQGDecorationTool.cxx
+++ b/PhysicsAnalysis/DerivationFramework/DerivationFrameworkMCTruth/src/TruthQGDecorationTool.cxx
@@ -39,11 +39,14 @@ StatusCode DerivationFramework::TruthQGDecorationTool::addBranches() const
   SG::AuxElement::Decorator<int> output_decorator(m_decOutput);
 
   for (auto ajet : *inputJets){
-    if (!ajet->isAvailable<int>("PartonTruthLabelID") ||
-        !ajet->isAvailable<int>("HadronConeExclTruthLabelID")){
-      ATH_MSG_ERROR("Did not have input decorations available");
+    if (!ajet->isAvailable<int>("PartonTruthLabelID") ){
+      ATH_MSG_ERROR("Did not have input PartonTruthLabelID decorations available");
       return StatusCode::FAILURE;
-    } // Now we have the input decorations          
+    }
+    else if (!ajet->isAvailable<int>("HadronConeExclTruthLabelID") ){
+      ATH_MSG_ERROR("Did not have input HadronConeExclTruthLabelID decorations available");
+      return StatusCode::FAILURE;
+    } // Now we have the input decorations
     /* Agreement from the HF-tagging and Jet/MET group:
         - If it is non-zero, use the label from the HF-tagging group (b, c, tau)
         - If it is zero, use the label from the Jet/MET group (q/g)
-- 
GitLab