Fix METMap for PHYS -> PHYSLITE production

Fix METMap for PHYS -> PHYSLITE production

PhysicsAnalysis/DerivationFramework/DerivationFrameworkJetEtMiss/CMakeLists.txt

-# Copyright (C) 2002-2020 CERN for the benefit of the ATLAS collaboration
+# Copyright (C) 2002-2023 CERN for the benefit of the ATLAS collaboration
 
 # Declare the package name:
 atlas_subdir( DerivationFrameworkJetEtMiss )
@@ -10,7 +10,7 @@ find_package( ROOT COMPONENTS Core Tree MathCore Hist RIO pthread )
 atlas_add_component( DerivationFrameworkJetEtMiss
                      src/*.cxx src/components/*.cxx
                      INCLUDE_DIRS ${ROOT_INCLUDE_DIRS}
-                     LINK_LIBRARIES ${ROOT_LIBRARIES} AsgTools AthLinks AthenaBaseComps DerivationFrameworkInterfaces FTagAnalysisInterfacesLib GaudiKernel InDetTrackSelectionToolLib JetAnalysisInterfacesLib JetInterface LumiBlockData PFlowUtilsLib ParticleJetToolsLib PathResolver StoreGateLib TrackVertexAssociationToolLib xAODCaloEvent xAODCore xAODEventInfo xAODJet xAODPFlow xAODTracking xAODTrigger xAODTruth )
+                     LINK_LIBRARIES ${ROOT_LIBRARIES} AsgTools AthLinks AthenaBaseComps DerivationFrameworkInterfaces FTagAnalysisInterfacesLib GaudiKernel InDetTrackSelectionToolLib JetAnalysisInterfacesLib JetInterface LumiBlockData PFlowUtilsLib ParticleJetToolsLib PathResolver StoreGateLib TrackVertexAssociationToolLib xAODCaloEvent xAODCore xAODEgamma xAODEventInfo xAODJet xAODMissingET xAODMuon xAODPFlow xAODTau xAODTracking xAODTrigger xAODTruth )
 
 # Install files from the package:
 atlas_install_python_modules( python/*.py POST_BUILD_CMD ${ATLAS_FLAKE8} )

PhysicsAnalysis/DerivationFramework/DerivationFrameworkJetEtMiss/python/METCommonConfig.py

-# Copyright (C) 2002-2022 CERN for the benefit of the ATLAS collaboration
+# Copyright (C) 2002-2023 CERN for the benefit of the ATLAS collaboration
 
 #********************************************************************
 # METCommonConfig.py
@@ -46,4 +46,12 @@ def METCustomVtxCfg(ConfigFlags, vxColl, jetColl, constituentColl):
     for assoc in cfg.assoclist:
         assoc.PrimVxColl = vxColl
 
-    return getAssocCA(cfg,METName='CustomJet')
\ No newline at end of file
+    return getAssocCA(cfg, METName='CustomJet')
+
+def METRemappingCfg(ConfigFlags):
+    from AthenaConfiguration.ComponentAccumulator import ComponentAccumulator, CompFactory
+
+    acc = ComponentAccumulator()
+    acc.addEventAlgo(CompFactory.DerivationFramework.METRemappingAlg('AnalysisMETRemappingAlg'))
+
+    return acc

PhysicsAnalysis/DerivationFramework/DerivationFrameworkJetEtMiss/python/MET_Baseline_AntiKt4EMPFlowCPContent.py

-# Copyright (C) 2002-2020 CERN for the benefit of the ATLAS collaboration
+# Copyright (C) 2002-2023 CERN for the benefit of the ATLAS collaboration
 
 MET_Baseline_AntiKt4EMPFlowCPContent  = [
 "AntiKt4EMPFlowJets",

PhysicsAnalysis/DerivationFramework/DerivationFrameworkJetEtMiss/src/METRemappingAlg.cxx

+///////////////////////// -*- C++ -*- /////////////////////////////
+
+/*
+  Copyright (C) 2002-2023 CERN for the benefit of the ATLAS collaboration
+*/
+
+#include "METRemappingAlg.h"
+
+#include <memory>
+
+namespace DerivationFramework {
+
+  METRemappingAlg::METRemappingAlg(const std::string& name, ISvcLocator* pSvcLocator) :
+    AthAlgorithm(name, pSvcLocator)
+  {
+
+  }
+
+  StatusCode METRemappingAlg::initialize()
+  {
+    ATH_MSG_VERBOSE("METRemappingAlg::initialize()");
+
+    ATH_CHECK( m_jetContKey.initialize() );
+    ATH_CHECK( m_photonContKey.initialize() );
+    ATH_CHECK( m_electronContKey.initialize() );
+    ATH_CHECK( m_muonContKey.initialize() );
+    ATH_CHECK( m_tauContKey.initialize() );
+
+    ATH_CHECK( m_inputMapKey.initialize() );
+    ATH_CHECK( m_inputCoreKey.initialize() );
+    ATH_CHECK( m_outputMapKey.initialize() );
+    ATH_CHECK( m_outputCoreKey.initialize() );
+
+    return StatusCode::SUCCESS;
+  }
+
+  StatusCode METRemappingAlg::execute()
+  {
+    ATH_MSG_VERBOSE("METRemappingAlg::execute()");
+
+    const EventContext& ctx = Gaudi::Hive::currentContext();
+
+    SG::ReadHandle<xAOD::JetContainer> jetContHandle(m_jetContKey, ctx);
+    if( !jetContHandle.isValid() ) {
+      ATH_MSG_ERROR("Unable to retrieve input jet container " << m_jetContKey.key());
+      return StatusCode::FAILURE;
+    }
+
+    // first iterate through the AnalysisJets container and populate a map
+    // that links original Jet objects to their calibrated counterparts
+    std::map<const xAOD::Jet*, ElementLink<xAOD::JetContainer> > jetLinkMap;
+    for( const xAOD::IParticle *j : *jetContHandle ) {
+      if( !m_accOriginalObject.isAvailable(*j) ) {
+        ATH_MSG_ERROR("originalObjectLink not available!");
+        return StatusCode::FAILURE;
+      }
+      const xAOD::IParticle *orig = *m_accOriginalObject(*j);
+      ElementLink<xAOD::JetContainer> link(*jetContHandle, j->index());
+      jetLinkMap.try_emplace(
+        static_cast<const xAOD::Jet*>(orig),
+        link
+      );
+    }
+
+    // repeat for Photon/Electron/Muon/Tau containers
+    linkMap_t objectLinkMap;
+    SG::ReadHandle<xAOD::PhotonContainer> photonContHandle(m_photonContKey, ctx);
+    ATH_CHECK( fillLinkMap(objectLinkMap, photonContHandle) );
+
+    SG::ReadHandle<xAOD::ElectronContainer> electronContHandle(m_electronContKey, ctx);
+    ATH_CHECK( fillLinkMap(objectLinkMap, electronContHandle) );
+
+    SG::ReadHandle<xAOD::MuonContainer> muonContHandle(m_muonContKey, ctx);
+    ATH_CHECK( fillLinkMap(objectLinkMap, muonContHandle) );
+
+    SG::ReadHandle<xAOD::TauJetContainer> tauContHandle(m_tauContKey, ctx);
+    ATH_CHECK( fillLinkMap(objectLinkMap, tauContHandle) );
+
+    // now retrieve and iterate through the METmap from PHYS and
+    // use its contents as a baseline to populate our own
+    SG::ReadHandle<xAOD::MissingETAssociationMap> inputMapHandle(m_inputMapKey, ctx);
+    if( !inputMapHandle.isValid() ) {
+      ATH_MSG_ERROR("Unable to retrieve input MissingETAssociationMap " << m_inputMapKey.key());
+      return StatusCode::FAILURE;
+    }
+
+    SG::WriteHandle<xAOD::MissingETAssociationMap> outputMapHandle = SG::makeHandle(m_outputMapKey, ctx);
+    ATH_CHECK( outputMapHandle.isValid() );
+    ATH_CHECK( outputMapHandle.record(
+                 std::make_unique<xAOD::MissingETAssociationMap>(),
+                 std::make_unique<xAOD::MissingETAuxAssociationMap>()
+                                      ));
+
+    const ElementLink<xAOD::IParticleContainer> invalidLink;
+    for( const xAOD::MissingETAssociation *el : *inputMapHandle ) {
+      // copy constructor creates a deep copy
+      auto assoc = outputMapHandle->push_back(new xAOD::MissingETAssociation(*el));
+
+      if( !assoc->isMisc() ) {
+        // check if the reference jet has a calibrated equivalent that should be linked to instead
+        std::map<const xAOD::Jet*, ElementLink<xAOD::JetContainer> >::const_iterator jet_it = jetLinkMap.find(assoc->refJet());
+        if( jet_it != jetLinkMap.end() ) {
+          // relink to calibrated jet
+          assoc->setJetLink(jet_it->second);
+          
+          // update objectLinks for this association
+          MissingETBase::Types::objlink_vector_t objectLinks;
+          for( const ElementLink<xAOD::IParticleContainer> &link : assoc->objectLinks() ) {
+            if( !link.isValid() ) {
+              objectLinks.push_back(invalidLink);
+              continue;
+            }
+
+            linkMap_t::const_iterator obj_it = objectLinkMap.find(*link);
+            if( obj_it != objectLinkMap.end() ) {
+              objectLinks.emplace_back(obj_it->second);
+            } else {
+              // objects that aren't found in the map were selected away,
+              // but we should leave an invalid link to maintain index order
+              objectLinks.push_back(invalidLink);
+            }
+          }
+          assoc->setObjectLinks(objectLinks);
+
+        } else { // jet_it == jetLinkMap.end()
+          // jet was selected away - this case should not happen, just give an error for now
+          ATH_MSG_ERROR("Jet not found!");
+          return StatusCode::FAILURE;
+        }
+      } else { // assoc->isMisc() == true
+        // update links in the misc association
+        MissingETBase::Types::objlink_vector_t miscObjectLinks;
+        for( const ElementLink<xAOD::IParticleContainer> &link : assoc->objectLinks() ) {
+          if( !link.isValid() ) {
+            miscObjectLinks.push_back(invalidLink);
+            continue;
+          }
+          
+          linkMap_t::const_iterator obj_it = objectLinkMap.find(*link);
+          if( obj_it != objectLinkMap.end() ) {
+            miscObjectLinks.emplace_back(obj_it->second);
+          } else {
+            miscObjectLinks.push_back(invalidLink);
+          }
+        }
+        assoc->setObjectLinks(miscObjectLinks);
+      }
+    } //> end loop over METmap
+
+    // copy over the MET core container
+    SG::ReadHandle<xAOD::MissingETContainer> inputCoreHandle(m_inputCoreKey, ctx);
+    if( !inputCoreHandle.isValid() ) {
+      ATH_MSG_ERROR("Unable to retrieve input MET core container " << m_inputCoreKey.key());
+      return StatusCode::FAILURE;
+    }
+
+    SG::WriteHandle<xAOD::MissingETContainer> outputCoreHandle = SG::makeHandle(m_outputCoreKey, ctx);
+    ATH_CHECK( outputCoreHandle.isValid() );
+    ATH_CHECK( outputCoreHandle.record(
+                 std::make_unique<xAOD::MissingETContainer>(),
+                 std::make_unique<xAOD::MissingETAuxContainer>()
+                                       ));
+
+    for( const xAOD::MissingET *el : *inputCoreHandle ) {
+      outputCoreHandle->push_back(new xAOD::MissingET(*el));
+    }
+
+    return StatusCode::SUCCESS;
+  }
+
+  template<typename handle_t> 
+  StatusCode METRemappingAlg::fillLinkMap(linkMap_t &map, handle_t &handle)
+  {
+    if( !handle.isValid() ) {
+      ATH_MSG_ERROR("Unable to retrieve " << handle.key());
+      return StatusCode::FAILURE;
+    }
+
+    for( const xAOD::IParticle *obj : *handle ) {
+      if( !m_accOriginalObject.isAvailable(*obj) ) {
+        ATH_MSG_ERROR("originalObjectLink not available!");
+        return StatusCode::FAILURE;
+      }
+      const xAOD::IParticle *orig = *m_accOriginalObject(*obj);
+      ElementLink<xAOD::IParticleContainer> link(*handle,obj->index());
+      map.try_emplace(orig, link);
+    }
+    return StatusCode::SUCCESS;
+  }
+
+} //> end namespace DerivationFramework

PhysicsAnalysis/DerivationFramework/DerivationFrameworkJetEtMiss/src/METRemappingAlg.h

+///////////////////////// -*- C++ -*- /////////////////////////////
+
+/*
+  Copyright (C) 2002-2023 CERN for the benefit of the ATLAS collaboration
+ */
+
+#ifndef DERIVATIONFRAMEWORK_METREMAPPINGALG_H
+#define DERIVATIONFRAMEWORK_METREMAPPINGALG_H
+
+#include <string>
+#include <vector>
+#include <map>
+
+#include "Gaudi/Property.h"
+#include "GaudiKernel/ToolHandle.h"
+#include "StoreGate/DataHandle.h"
+#include "AthenaBaseComps/AthAlgorithm.h"
+#include "AthLinks/ElementLink.h"
+
+#include "xAODBase/IParticle.h"
+#include "xAODBase/IParticleContainer.h"
+#include "xAODEgamma/ElectronContainer.h"
+#include "xAODEgamma/PhotonContainer.h"
+#include "xAODMuon/MuonContainer.h"
+#include "xAODTau/TauJetContainer.h"
+#include "xAODMissingET/MissingETAssociationMap.h"
+#include "xAODMissingET/MissingETAuxAssociationMap.h"
+#include "xAODMissingET/MissingETContainer.h"
+#include "xAODMissingET/MissingETAuxContainer.h"
+
+namespace DerivationFramework {
+  class METRemappingAlg : public AthAlgorithm {
+  public:
+    METRemappingAlg(const std::string& name, ISvcLocator* pSvcLocator);
+    virtual ~METRemappingAlg() = default;
+
+    virtual StatusCode initialize() override;
+    virtual StatusCode execute() override;
+    
+  private:
+    typedef std::map<const xAOD::IParticle*, ElementLink<xAOD::IParticleContainer>> linkMap_t;
+    template<typename handle_t> StatusCode fillLinkMap(linkMap_t &map, handle_t &handle);
+
+    SG::ReadHandleKey<xAOD::JetContainer> m_jetContKey{this, "JetCollectionKey", "AnalysisJets", "SG key for the analysis jets collection"};
+    SG::ReadHandleKey<xAOD::PhotonContainer> m_photonContKey{this, "PhotonCollectionKey", "AnalysisPhotons", "SG key for the analysis photons collection"};
+    SG::ReadHandleKey<xAOD::ElectronContainer> m_electronContKey{this, "ElectronCollectionKey", "AnalysisElectrons", "SG key for the analysis electrons collection"};
+    SG::ReadHandleKey<xAOD::MuonContainer> m_muonContKey{this, "MuonCollectionKey", "AnalysisMuons", "SG key for the analysis muons collection"};
+    SG::ReadHandleKey<xAOD::TauJetContainer> m_tauContKey{this, "TauCollectionKey", "AnalysisTauJets", "SG key for the analysis tau jets collection"};
+    SG::ReadHandleKey<xAOD::MissingETAssociationMap> m_inputMapKey{this, "AssociationInputKey", "METAssoc_AntiKt4EMPFlow", "SG key for the input MissingETAssociationMap"};
+    SG::ReadHandleKey<xAOD::MissingETContainer> m_inputCoreKey{this, "METCoreInputKey", "MET_Core_AntiKt4EMPFlow", "SG key for the input MET core container"};
+    SG::WriteHandleKey<xAOD::MissingETAssociationMap> m_outputMapKey{this, "AssociationOutputKey", "METAssoc_AnalysisMET", "SG key for the output MissingETAssociationMap"};
+    SG::WriteHandleKey<xAOD::MissingETContainer> m_outputCoreKey{this, "METCoreOutputKey", "MET_Core_AnalysisMET", "SG key for the output MET core container"};
+
+    const SG::AuxElement::ConstAccessor< ElementLink<xAOD::IParticleContainer> > m_accOriginalObject{"originalObjectLink"};
+
+  }; //> end class METRemappingAlg
+} //> end namespace DerivationFramework
+
+#endif //> !DERIVATIONFRAMEWORK_METREMAPPINGALG_H

PhysicsAnalysis/DerivationFramework/DerivationFrameworkJetEtMiss/src/components/DerivationFrameworkJetEtMiss_entries.cxx

 #include "../PFlowAugmentationTool.h"
+#include "../METRemappingAlg.h"
 #include "../METTriggerAugmentationTool.h"
 #include "../ViewContainerThinning.h"
 #include "../JetExternalAssocTool.h"
@@ -9,6 +10,7 @@
 using namespace DerivationFramework;
 
 DECLARE_COMPONENT( PFlowAugmentationTool )
+DECLARE_COMPONENT( METRemappingAlg )
 DECLARE_COMPONENT( METTriggerAugmentationTool )
 DECLARE_COMPONENT( ViewContainerThinning )
 DECLARE_COMPONENT( JetExternalAssocTool )

PhysicsAnalysis/DerivationFramework/DerivationFrameworkPhys/python/PHYSLITE.py

@@ -132,7 +132,6 @@ def PHYSLITEKernelCfg(ConfigFlags, name='PHYSLITEKernel', **kwargs):
         acc.addEventAlgo(element)
     
     # Build MET from our analysis objects
-    # Currently only works for the AOD->PHYSLITE workflow
     if 'StreamAOD' in ConfigFlags.Input.ProcessingTags:
         from METReconstruction.METAssocCfg import AssocConfig, METAssocConfig
         from METReconstruction.METAssociatorCfg import getAssocCA
@@ -149,6 +148,11 @@ def PHYSLITEKernelCfg(ConfigFlags, name='PHYSLITEKernel', **kwargs):
                                       usePFOLinks=True)
         components_PHYSLITE_cfg = getAssocCA(PHYSLITE_cfg,METName='AnalysisMET')
         acc.merge(components_PHYSLITE_cfg)
+    elif 'StreamDAOD_PHYS' in ConfigFlags.Input.ProcessingTags:
+        from DerivationFrameworkJetEtMiss.METCommonConfig import METRemappingCfg
+
+        METRemap_cfg = METRemappingCfg(ConfigFlags)
+        acc.merge(METRemap_cfg)
 
     # The derivation kernel itself
     DerivationKernel = CompFactory.DerivationFramework.DerivationKernel