Support origin correction in derivations

PhysicsAnalysis/DerivationFramework/DerivationFrameworkJetEtMiss/python/ExtendedJetCommon.py

@@ -631,6 +631,24 @@ def applyBTaggingAugmentation(jetalg,algname='default',sequence=DerivationFramew
     extjetlog.info('ExtendedJetCommon: Applying b-tagging working points for jet collection: '+inJets)
     applyJetAugmentation(jetalg,algname,sequence,jetaugtool)
 
+def addOriginCorrection(jetalg, sequence, algname,vertexPrefix):
+    jetaugtool = getJetAugmentationTool(jetalg,'_OriginCorr'+vertexPrefix)
+    if(jetaugtool is None):
+        extjetlog.warning('*** addOriginCorrection called but corresponding augmentation tool does not exist! ***')
+
+    origincorrectiontoolname = 'DFOriginCorrection'+vertexPrefix+'_'+jetalg
+    jetaugtool.MomentPrefix = vertexPrefix+'_'
+    from AthenaCommon.AppMgr import ToolSvc
+    if hasattr(ToolSvc,origincorrectiontoolname):
+        jetaugtool.JetOriginCorrectionTool = getattr(ToolSvc,origincorrectiontoolname)
+    else:
+        origincorrectiontool = CfgMgr.JetOriginCorrectionTool(origincorrectiontoolname, JetContainer=jetalg+'Jets', VertexContainer=vertexPrefix+'PrimaryVertices',OriginCorrectedName=vertexPrefix+'_JetOriginConstitScaleMomentum',ForceEMScale=True)
+        ToolSvc += origincorrectiontool
+        jetaugtool.JetOriginCorrectionTool = origincorrectiontool
+
+    extjetlog.info('ExtendedJetCommon: Adding OriginCorrection for jet collection: '+jetalg)
+    applyJetAugmentation(jetalg,algname,sequence,jetaugtool)
+
 #################################################################
 ### Schedule Q/G-tagging decorations ### QGTaggerTool ##### 
 #################################################################

PhysicsAnalysis/DerivationFramework/DerivationFrameworkJetEtMiss/python/JetCommon.py

@@ -66,7 +66,8 @@ def defineEDAlg(R=0.4, inputtype="LCTopo"):
                     "LCTopoOrigin" : jtm.lcoriginget,
                     "EMTopoOrigin" : jtm.emoriginget,
                     "EMPFlow": jtm.empflowget,
-                    "EMPFlowFE": jtm.empflowget_fe
+                    "EMPFlowFE": jtm.empflowget_fe,
+                    "PFlowCustomVtx": jtm.pflowcustomvtxget
                     }[inputtype]
 
     t=configEventDensityTool("EDTool"+str(int(R*10))+inputtype,

PhysicsAnalysis/DerivationFramework/DerivationFrameworkJetEtMiss/src/JetAugmentationTool.cxx

@@ -27,6 +27,7 @@ namespace DerivationFramework {
     m_dobtag(false),
     m_jetTrackSumMomentsTool(""),
     m_decoratetracksum(false),
+    m_decorateorigincorrection(false),
     m_jetPtAssociationTool(""),
     m_decorateptassociation(false),
     m_trkSelectionTool(""),
@@ -222,6 +223,24 @@ namespace DerivationFramework {
       ATH_CHECK(m_truthLabel_truthJetMass_key.initialize());
     }
 
+    if(!m_jetOriginCorrectionTool.empty()) {
+      CHECK(m_jetOriginCorrectionTool.retrieve());
+      ATH_MSG_DEBUG("Augmenting jets with origin corrections \"" << m_momentPrefix << "Origin\"");
+      m_decorateorigincorrection = true;
+      m_origincorrection_key = m_containerName + "." + m_momentPrefix + "OriginVertex";
+      m_originpt_key         = m_containerName + "." + m_momentPrefix + "pt";
+      m_origineta_key        = m_containerName + "." + m_momentPrefix + "eta";
+      m_originphi_key        = m_containerName + "." + m_momentPrefix + "phi";
+      m_originm_key          = m_containerName + "." + m_momentPrefix + "m";
+
+      ATH_CHECK(m_origincorrection_key.initialize());
+      ATH_CHECK(m_originpt_key.initialize());
+      ATH_CHECK(m_origineta_key.initialize());
+      ATH_CHECK(m_originphi_key.initialize());
+      ATH_CHECK(m_originm_key.initialize());
+    }
+
+
     return StatusCode::SUCCESS;
   }
     
@@ -287,6 +306,13 @@ namespace DerivationFramework {
 	}
     }
 
+    if(m_decorateorigincorrection){
+      if(m_jetOriginCorrectionTool->modify(*jets_copy).isFailure()){
+        ATH_MSG_WARNING("Problem applying the origin correction tool");
+        return StatusCode::FAILURE;
+      }
+    }
+
     // Check if GhostTruthAssociation decorations already exist for first jet, and if so skip them
     bool isMissingPtAssociation = true;
     if( !m_decorateptassociation || jets_copy->size() == 0 || m_acc_GhostTruthAssociationFraction->isAvailable(*jets_copy->at(0)) ) {
@@ -385,6 +411,20 @@ namespace DerivationFramework {
 	ATH_MSG_VERBOSE("TrackSumPt: "   << (*m_acc_tracksummass)(jet_orig) );
       }
 
+      if(m_decorateorigincorrection) {
+        SG::WriteDecorHandle<xAOD::JetContainer, float> originPtDec(m_originpt_key);
+        SG::WriteDecorHandle<xAOD::JetContainer, float> originEtaDec(m_origineta_key);
+        SG::WriteDecorHandle<xAOD::JetContainer, float> originPhiDec(m_originphi_key);
+        SG::WriteDecorHandle<xAOD::JetContainer, float> originMDec(m_originm_key);
+        SG::WriteDecorHandle<xAOD::JetContainer, ElementLink<xAOD::VertexContainer> > originVxDec(m_origincorrection_key);
+        originPtDec(jet_orig)  = jet->pt();
+        originEtaDec(jet_orig) = jet->eta();
+        originPhiDec(jet_orig) = jet->phi();
+        originMDec(jet_orig)   = jet->m();
+        originVxDec(jet_orig)  = jet->getAttribute<ElementLink<xAOD::VertexContainer> >("OriginVertex");
+      }
+
+
       if(m_decorateptassociation && isMissingPtAssociation){
 
 	SG::WriteDecorHandle<xAOD::JetContainer, float> ghostTruthAssocFrac_handle(m_GhostTruthAssociationFraction_key);

PhysicsAnalysis/DerivationFramework/DerivationFrameworkJetEtMiss/src/JetAugmentationTool.h

@@ -90,6 +90,15 @@ namespace DerivationFramework {
     SG::WriteDecorHandleKey<xAOD::JetContainer> m_tracksummass_key {this,"TrackSumMassKey", "", "Decoration for mass calculated from associated tracks"};
     SG::WriteDecorHandleKey<xAOD::JetContainer> m_tracksumpt_key {this,"TrackSumPtKey", "","Decoration for pt calculated from associated tracks"};
 
+    // Origin correction
+    ToolHandle<IJetModifier> m_jetOriginCorrectionTool{this, "JetOriginCorrectionTool", "", "Origin correction tool"};
+    bool m_decorateorigincorrection;
+    SG::WriteDecorHandleKey<xAOD::JetContainer> m_origincorrection_key{this, "OriginVertex", "", "Origin vertex (Autoconfigured)"};
+    SG::WriteDecorHandleKey<xAOD::JetContainer> m_originpt_key{this, "OriginPt", "", "Origin corrected pt (Autoconfigured)"};
+    SG::WriteDecorHandleKey<xAOD::JetContainer> m_origineta_key{this, "OriginEta", "", "Origin corrected eta (Autoconfigured)"};
+    SG::WriteDecorHandleKey<xAOD::JetContainer> m_originphi_key{this, "OriginPhi", "", "Origin corrected phi (Autoconfigured)"};
+    SG::WriteDecorHandleKey<xAOD::JetContainer> m_originm_key{this, "OriginM", "", "Origin corrected mass (Autoconfigured)"};
+
     // GhostTruthAssociation for derivations, @author jeff.dandoy@cern.ch
     ToolHandle<IJetModifier> m_jetPtAssociationTool;
     bool m_decorateptassociation;

Reconstruction/Jet/JetMomentTools/JetMomentTools/JetOriginCorrectionTool.h

@@ -51,6 +51,7 @@ public:
   Gaudi::Property<std::string> m_correctionName{this, "OriginCorrectedName", "JetOriginConstitScaleMomentum", "Origin corrected name"};
   Gaudi::Property<std::string> m_jetContainerName{this, "JetContainer", "", "SG key for the input jet container"};
   Gaudi::Property<bool> m_onlyAssignPV{this, "OnlyAssignPV", false, "Only write out PV information"};
+  Gaudi::Property<bool> m_forceEMScale{this, "ForceEMScale", false, "Force apply origin correction at EM scale"};
 
 private:
 

Reconstruction/Jet/JetMomentTools/Root/JetOriginCorrectionTool.cxx

@@ -122,7 +122,9 @@ StatusCode JetOriginCorrectionTool::decorate(const xAOD::JetContainer& jetCont)
     ATH_MSG_DEBUG("                     jet pT: "<< jet->pt());
 
     if(!m_onlyAssignPV) {
-      xAOD::JetFourMom_t fv = jet::clusterOriginCorrection(*jet,*vx);
+      xAOD::CaloCluster::State state = xAOD::CaloCluster::UNKNOWN;
+      if(m_forceEMScale) state = xAOD::CaloCluster::UNCALIBRATED;
+      xAOD::JetFourMom_t fv = jet::clusterOriginCorrection(*jet,*vx,state);
       ATH_MSG_DEBUG("  " <<  m_correctionName << " pT: " << fv.pt());
       scaleMomentumPtHandle(*jet) = fv.pt();
       scaleMomentumPhiHandle(*jet) = fv.phi();

Reconstruction/Jet/JetRec/python/JetRecStandardTools.py

@@ -387,6 +387,15 @@ jtm += PseudoJetAlgorithm(
   SkipNegativeEnergy = True,
 )
 
+# EM-scale pflow with custom selection for the primary vertex 
+jtm += PseudoJetAlgorithm(
+  "pflowcustomvtxget",
+  Label = "PFlowCustomVtx",
+  InputContainer = "CustomVtxParticleFlowObjects",
+  OutputContainer = "PseudoJetPFlowCustomVtx",
+  SkipNegativeEnergy = True,
+)
+
 # AntiKt2 track jets.
 jtm += PseudoJetAlgorithm(
   "gakt2trackget", # give a unique name

Reconstruction/Jet/JetUtils/JetUtils/JetOriginHelpers.h

@@ -21,6 +21,7 @@
 ////////////////////////////////////////////////////////////
 #include "xAODJet/Jet.h"
 #include "xAODTracking/VertexFwd.h"
+#include "xAODCaloEvent/CaloClusterContainer.h"
 
 namespace jet {
 
@@ -32,6 +33,7 @@ namespace jet {
   /// The cluster constituents of jet are retrieved, set at the scale they were used to find
   /// jet, corrected according to vx. The 4-vector sum of this corrected cluster is returned
   xAOD::JetFourMom_t clusterOriginCorrection(const xAOD::Jet& jet, const xAOD::Vertex & vx);
+  xAOD::JetFourMom_t clusterOriginCorrection(const xAOD::Jet& jet, const xAOD::Vertex & vx, const xAOD::CaloCluster::State state = xAOD::CaloCluster::UNKNOWN);
 
 }
 

Reconstruction/Jet/JetUtils/Root/JetOriginHelpers.cxx

@@ -9,10 +9,12 @@
 
 namespace jet {
 
-  xAOD::JetFourMom_t clusterOriginCorrection(const xAOD::Jet& jet, const xAOD::Vertex & vx){
+  xAOD::JetFourMom_t clusterOriginCorrection(const xAOD::Jet& jet, const xAOD::Vertex & vx, const xAOD::CaloCluster::State state){
     
 
-    xAOD::CaloCluster::State  constitScale = (xAOD::CaloCluster::State) jet.getConstituentsSignalState();
+    xAOD::CaloCluster::State constitScale = state;
+    if(constitScale==xAOD::CaloCluster::UNKNOWN) constitScale = (xAOD::CaloCluster::State) jet.getConstituentsSignalState();
+
     int numC = jet.numConstituents();
     const Amg::Vector3D& vxpos= vx.position();