METAssociator.cxx

Go to the documentation of this file.

 
/*
  Copyright (C) 2002-2025 CERN for the benefit of the ATLAS collaboration
*/
 
// METAssociator.cxx
// Implementation for class METAssociator
//
// This is the base class for tools that construct MET terms
// from other object collections.
//
//  * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * * *
//
// Author: P Loch, S Resconi, TJ Khoo, AS Mete
 
// METReconstruction includes
#include "METReconstruction/METAssociator.h"
#include "METReconstruction/METEgammaAssociator.h"
#include "xAODMissingET/MissingETComposition.h"
#include "xAODMissingET/MissingETContainer.h"
#include "xAODMissingET/MissingETAssociationMap.h"
#include "AthContainers/ConstDataVector.h"
 
// Tracking EDM
#include "xAODTracking/Vertex.h"
#include "xAODTracking/TrackParticle.h"
 
// Track errors
#include "EventPrimitives/EventPrimitivesHelpers.h"
 
// For DeltaR
#include "FourMomUtils/xAODP4Helpers.h"
 
#include "AsgDataHandles/WriteDecorHandle.h"
 
#include "GaudiKernel/SystemOfUnits.h"
 
#include "TRandom3.h"
 
 
namespace met {
 
  using namespace xAOD;

  // Public methods:
 
  // Constructors
  METAssociator::METAssociator(const std::string& name) :
    AsgTool(name)
  {
  }
 
  // Destructor
  METAssociator::~METAssociator()
  = default;
 
  // Athena algtool's Hooks
  StatusCode METAssociator::initialize()
  {
    ATH_MSG_DEBUG ("Initializing " << name() << "...");
 
    ATH_CHECK( m_trkseltool.retrieve() );
    ATH_CHECK(m_trkIsolationTool.retrieve());
    ATH_CHECK(m_caloIsolationTool.retrieve());
 
    if(m_clcollKey.key() == "CaloCalTopoClusters") {
      if(m_useModifiedClus) {
        ATH_MSG_ERROR("Configured to use modified topocluster collection but \"CaloCalTopoClusters\" collection specified!");
        return StatusCode::FAILURE;
      } else {
        ATH_MSG_INFO("Configured to use standard topocluster collection.");
      }
    } else {
      if(m_useModifiedClus) {
        ATH_MSG_INFO("Configured to use modified topocluster collection \"" << m_clcollKey.key() << "\".");
      } else {
        ATH_MSG_ERROR("Configured to use topocluster collection \"" << m_clcollKey.key() << "\", but modified clusters flag not set!");
        return StatusCode::FAILURE;
      }
    }
    //initialise read handle keys
    ATH_CHECK( m_pvcollKey.initialize(m_useTracks));
    ATH_CHECK( m_trkcollKey.initialize(m_useTracks));
 
    ATH_CHECK(m_fecollKey.initialize(m_pflow && !m_fecollKey.key().empty()));
    ATH_CHECK( m_pfcollKey.initialize(m_pflow && m_fecollKey.key().empty()));
    if(m_pflow){
      if(!m_fecollKey.key().empty()){
        ATH_MSG_INFO("Configured to use FlowElement collection \"" << m_fecollKey.key() << "\".");
      }
      else{
        ATH_MSG_INFO("Configured to use PFlow collection \"" << m_pfcollKey.key() << "\".");
        if(m_pfcollKey.key() == "JetETMissParticleFlowObjects") {
          ATH_MSG_ERROR("Configured to use standard pflow collection \"" << m_pfcollKey.key() << "\".");
          ATH_MSG_ERROR("This is no longer supported -- please use the CHSParticleFlowObjects collection, which has the four-vector corrections built in.");
          return StatusCode::FAILURE;
        }
      }
    }
 
    ATH_CHECK( m_clcollKey.initialize(!m_skipconst || m_forcoll.value().empty()));
 
    std::string hybridname = "Etmiss";
    hybridname += m_clcollKey.key();
    hybridname += m_foreta.value();
    hybridname += m_forcoll.value();
    ATH_CHECK( m_hybridContKey.assign(hybridname)); // FIXME Keys should be properly assigned during the configuration
    ATH_CHECK( m_hybridContKey.initialize(m_skipconst && !m_forcoll.value().empty()));
 
    ATH_CHECK(m_UEcorrPtDecorKey.initialize(SG::AllowEmpty));
    if (!m_recoil && !m_UEcorrPtDecorKey.key().empty()) {
      renounce(m_UEcorrPtDecorKey);
    }
    return StatusCode::SUCCESS;
  }
 
  StatusCode METAssociator::execute(xAOD::MissingETContainer* metCont, xAOD::MissingETAssociationMap* metMap, const EventContext& ctx) const
  {
    ATH_MSG_DEBUG ("In execute: " << name() << "...");
    if(!metCont) {
      ATH_MSG_WARNING("Invalid pointer to MissingETContainer supplied! Abort.");
      return StatusCode::FAILURE;
    }
 
    if(!metMap) {
      ATH_MSG_WARNING("Invalid pointer to MissingETAssociationMap supplied! Abort.");
      return StatusCode::FAILURE;
    }
    if(m_pflow && !m_useTracks ){
      ATH_MSG_WARNING("Attempting to build PFlow MET without a track collection.");
      return StatusCode::FAILURE;
    }
 
    return this->executeTool(metCont, metMap, ctx);
  }
 
  StatusCode METAssociator::retrieveConstituents(met::METAssociator::ConstitHolder& constits, const EventContext& ctx) const
  {
    ATH_MSG_DEBUG ("In execute: " << name() << "...");
    if (!m_skipconst || m_forcoll.value().empty()) { // FIXME m_clcollKey.value().empty() ???
 
      SG::ReadHandle<IParticleContainer> topoclusterCont(m_clcollKey, ctx);
      if (!topoclusterCont.isValid()) {
        ATH_MSG_WARNING("Unable to retrieve topocluster container " << m_clcollKey.key() << " for overlap removal");
        return StatusCode::FAILURE;
      }
      constits.tcCont=topoclusterCont.cptr();
      ATH_MSG_DEBUG("Successfully retrieved topocluster collection");
    } else {
      std::string hybridname = "Etmiss";
      hybridname += m_clcollKey.key();
      hybridname += m_foreta.value();
      hybridname += m_forcoll.value();
 
      SG::ReadHandle<IParticleContainer> hybridCont(m_hybridContKey, ctx);
      if( hybridCont.isValid()) {
        constits.tcCont=hybridCont.cptr();
      } else {
        ATH_MSG_WARNING("Trying to do something currently unsupported- lets abort");
        return StatusCode::FAILURE;
        // Trying to do this using write handles (need to get some input here)
        /*std::unique_ptr<ConstDataVector<IParticleContainer>> hybridCont = std::make_unique<ConstDataVector<IParticleContainer>>();
        SG::WriteHandle<ConstDataVector<IParticleContainer>> hybridContHandle(hybridname, ctx);
 
        StatusCode sc = hybridContHandle.record(std::make_unique<ConstDataVector<IParticleContainer>>(*hybridCont));
 
        if (sc.isFailure()) {
          ATH_MSG_WARNING("Unable to record container");
         return StatusCode::SUCCESS;
 
        }*/
 
        /*SG::ReadHandle<IParticleContainer> centCont(m_clcoll, ctx);
        if (!centCont.isValid()) {
          ATH_MSG_WARNING("Unable to retrieve central container " << m_clcoll << " for overlap removal");
          return StatusCode::FAILURE;
        }
 
        SG::ReadHandle<IParticleContainer> forCont(m_forcoll, ctx);
        if (!forCont.isValid()) {
          ATH_MSG_WARNING("Unable to retrieve forward container " << m_forcoll << " for overlap removal");
          return StatusCode::FAILURE;
        }
        ConstDataVector<IParticleContainer> *hybridCont = new ConstDataVector<IParticleContainer>(SG::VIEW_ELEMENTS);
 
        const IParticleContainer* centCont=0;
        if( evtStore()->retrieve(centCont, m_clcoll).isFailure() ) {
          ATH_MSG_WARNING("Unable to retrieve central container " << m_clcoll << " for overlap removal");
          return StatusCode::FAILURE;
        }
 
        const IParticleContainer* forCont=0;
        if( evtStore()->retrieve(forCont, m_forcoll).isFailure() ) {
          ATH_MSG_WARNING("Unable to retrieve forward container " << m_forcoll << " for overlap removal");
          return StatusCode::FAILURE;
        }
 
        for(const auto clus : *centCont) if (fabs(clus->eta())<m_foreta) hybridCont->push_back(clus);
        for(const auto clus : *forCont) if (fabs(clus->eta())>=m_foreta) hybridCont->push_back(clus);
        ATH_CHECK( evtStore()->record(hybridCont,hybridname));
        constits.tcCont = hybridCont->asDataVector();
        */
      }
    }
 
    if( !m_useTracks){
      //if you want to skip tracks, set the track collection empty manually
      ATH_MSG_DEBUG("Skipping tracks");
    }else{
      SG::ReadHandle<VertexContainer> vxCont(m_pvcollKey, ctx);
      if (!vxCont.isValid()) {
        ATH_MSG_WARNING("Unable to retrieve primary vertex container " << m_pvcollKey.key());
        //this is actually really bad.  If it's empty that's okay
        return StatusCode::FAILURE;
      }
 
      ATH_MSG_DEBUG("Successfully retrieved primary vertex container");
      ATH_MSG_DEBUG("Container holds " << vxCont->size() << " vertices");
 
      for(const auto *const vx : *vxCont) {
        ATH_MSG_VERBOSE( "Testing vertex " << vx->index() );
        if(vx->vertexType()==VxType::PriVtx)
          {constits.pv = vx; break;}
      }
      if(!constits.pv) {
        ATH_MSG_DEBUG("Failed to find primary vertex! Reject all tracks.");
      } else {
        ATH_MSG_VERBOSE("Primary vertex has z = " << constits.pv->z());
      }
 
      constits.trkCont=nullptr;
      ATH_MSG_DEBUG("Retrieving Track collection " << m_trkcollKey.key());
      SG::ReadHandle<TrackParticleContainer> trCont(m_trkcollKey, ctx);
      if (!trCont.isValid()) {
        ATH_MSG_WARNING("Unable to retrieve track particle container");
        return StatusCode::FAILURE;
      }
      constits.trkCont=trCont.cptr();
 
      if(m_pflow){
        if(!m_fecollKey.key().empty()){
          ATH_MSG_DEBUG("Retrieving FlowElement collection " << m_fecollKey.key());
          constits.feCont = nullptr;
          SG::ReadHandle<xAOD::FlowElementContainer> feCont(m_fecollKey, ctx);
          if (!feCont.isValid()) {
            ATH_MSG_ERROR("Unable to retrieve FlowElement container "<< m_fecollKey.key());
            return StatusCode::FAILURE;
          }
          constits.feCont=feCont.cptr();
        }
        else{
          ATH_MSG_DEBUG("Retrieving PFlow collection " << m_pfcollKey.key());
          constits.pfoCont = nullptr;
          SG::ReadHandle<PFOContainer> pfCont(m_pfcollKey, ctx);
          if (!pfCont.isValid()) {
            ATH_MSG_WARNING("Unable to PFlow object container");
            return StatusCode::FAILURE;
          }
          constits.pfoCont=pfCont.cptr();
        }
      }//pflow
    }//retrieve track/pfo containers
 
    return StatusCode::SUCCESS;
  }

  // Protected methods:
 
  StatusCode METAssociator::fillAssocMap(xAOD::MissingETAssociationMap* metMap,
                                         const xAOD::IParticleContainer* hardObjs, const EventContext& ctx) const
  {
    ConstitHolder constits;
 
    if (retrieveConstituents(constits, ctx).isFailure()) {
      ATH_MSG_DEBUG("Unable to retrieve constituent containers");
      return StatusCode::FAILURE;
    }
 
    SG::WriteDecorHandle<xAOD::IParticleContainer, float> dec_UEcorr (m_UEcorrPtDecorKey, ctx);
    std::vector<const IParticle*> constlist;
    constlist.reserve(20);
    std::vector<const IParticle*> hardObjs_tmp;
    for(const auto *const obj : *hardObjs) {
      hardObjs_tmp.push_back(obj);
    }
    std::sort(hardObjs_tmp.begin(),hardObjs_tmp.end(),greaterPt);
 
    for(const auto& obj : hardObjs_tmp) {
      if(obj->pt()<4e3 && obj->type()!=xAOD::Type::Muon){
        if (m_pflow && !m_fecollKey.key().empty() && m_useTracks && m_recoil) { dec_UEcorr(*obj) = 0.f; }
        continue;
      }
      constlist.clear();
      ATH_MSG_VERBOSE( "Object type, pt, eta, phi = " << obj->type() << ", " << obj->pt() << ", " << obj->eta() << "," << obj->phi() );
      if(m_pflow){
        if(!m_fecollKey.key().empty()){
          if(!m_useTracks){
            ATH_MSG_ERROR("Attempting to build FlowElement MET without a track collection.");
            return StatusCode::FAILURE;
          }
          std::map<const IParticle*, MissingETBase::Types::constvec_t> momentumOverride;
          if(m_recoil){ // HR part:
            float UEcorr_Pt = 0.; // Underlying event correction for HR
            ATH_CHECK(extractFEHR(obj,hardObjs_tmp,constlist,constits,momentumOverride, UEcorr_Pt));
            ATH_MSG_DEBUG("Energy correction is: " << UEcorr_Pt);
            dec_UEcorr(*obj) = UEcorr_Pt;
          }
          else{ // MET part:
            ATH_CHECK( this->extractFE(obj, constlist, constits, momentumOverride, ctx) );
          }
          MissingETComposition::insert(metMap, obj, constlist, momentumOverride);
        }
        else{
          // Old PFO EDM
          if(!m_useTracks){
            ATH_MSG_DEBUG("Attempting to build PFlow without a track collection.");
            return StatusCode::FAILURE;
          }else{
            std::map<const IParticle*,MissingETBase::Types::constvec_t> momentumOverride;
            ATH_CHECK( this->extractPFO(obj,constlist,constits,momentumOverride, ctx) );
            MissingETComposition::insert(metMap,obj,constlist,momentumOverride);
          }
        }
      } else {
        std::vector<const IParticle*> tclist;
        tclist.reserve(20);
        ATH_CHECK( this->extractTopoClusters(obj,tclist,constits, ctx) );
        if(m_useModifiedClus) {
          for(const auto& cl : tclist) {
            // use index-parallelism to identify shallow copied constituents
            constlist.push_back((*constits.tcCont)[cl->index()]);
          }
        } else {
          constlist = tclist;
        }
        if(m_useTracks) ATH_CHECK( this->extractTracks(obj,constlist,constits) );
        MissingETComposition::insert(metMap,obj,constlist);
      }
    }
    return StatusCode::SUCCESS;
  }
 
  // Accept Track
  bool METAssociator::acceptTrack(const xAOD::TrackParticle* trk, const xAOD::Vertex* vx) const
  {
 
    if (!vx) return false;//in events with no pv, we will just reject all tracks, and therefore build only the calo MET
    return static_cast<bool>  (m_trkseltool->accept( *trk, vx ));
 }
 
 
  bool METAssociator::isGoodEoverP(const xAOD::TrackParticle* trk) const
  {
 
    if( (fabs(trk->eta())<1.5 && trk->pt()>m_cenTrackPtThr) ||
        (fabs(trk->eta())>=1.5 && trk->pt()>m_forTrackPtThr) ) {
 
      // Get relative error on qoverp
      float Rerr = Amg::error(trk->definingParametersCovMatrix(),4)/fabs(trk->qOverP());
      ATH_MSG_VERBOSE( "Track momentum error (%): " << Rerr*100 );
 
      // first compute track and calo isolation variables
      float ptcone20 = 0., isolfrac = 0., etcone10 = 0., EoverP = 0.;
      // ptcone
      TrackIsolation trkIsoResult;
      std::vector<Iso::IsolationType> trkIsoCones;
      trkIsoCones.push_back(xAOD::Iso::IsolationType::ptcone20);
      xAOD::TrackCorrection trkIsoCorr;
      trkIsoCorr.trackbitset.set(xAOD::Iso::IsolationTrackCorrection::coreTrackPtr);
      m_trkIsolationTool->trackIsolation(trkIsoResult,
                                         *trk,
                                         trkIsoCones,
                                         trkIsoCorr);
      ptcone20 = !trkIsoResult.ptcones.empty() ? trkIsoResult.ptcones[0] : 0;
      isolfrac = ptcone20/trk->pt();
      // etcone
      CaloIsolation caloIsoResult;
      std::vector<Iso::IsolationType> caloIsoCones;
      // We can't actually configure the tool to give etcone10, so instead we have to compute etcone20,
      // applying the core cone correction.
      // Then, we retrieve the correction value, which is etcone10, rather than the isolation value
      caloIsoCones.push_back(xAOD::Iso::IsolationType::etcone20);
      xAOD::CaloCorrection caloIsoCorr_coreCone;
      caloIsoCorr_coreCone.calobitset.set(xAOD::Iso::IsolationCaloCorrection::coreCone); // this is etcone10
      m_caloIsolationTool->caloTopoClusterIsolation(caloIsoResult,
                                                    *trk,
                                                    caloIsoCones,
                                                    caloIsoCorr_coreCone);
      if(!caloIsoResult.etcones.empty()) {
        // retrieve the correction value for the core cone
        etcone10 = caloIsoResult.coreCorrections[xAOD::Iso::IsolationCaloCorrection::coreCone][xAOD::Iso::IsolationCorrectionParameter::coreEnergy];
      } else {
        ATH_MSG_WARNING("isGoodEoverP: Failed to retrieve the isolation core correction (etcone10)! Setting etcone10=0");
        etcone10 = 0.;
      }
      EoverP   =  etcone10/trk->pt();
      ATH_MSG_VERBOSE( "Track isolation fraction: " << isolfrac );
      ATH_MSG_VERBOSE( "Track E/P = " << EoverP );
 
      if(isolfrac<0.1) {
            // isolated track cuts
            if(Rerr>0.4) return false;
            else if (EoverP<0.65 && ((EoverP>0.1 && Rerr>0.05) || Rerr>0.1)) return false;
      } else {
            // non-isolated track cuts
            float trkptsum = ptcone20+trk->pt();
            if(etcone10/trkptsum<0.6 && trk->pt()/trkptsum>0.6) return false;
      }
    }
    return true;
  }
 
  StatusCode METAssociator::GetUEcorr(const met::METAssociator::ConstitHolder& constits,  // all PFOs
                                      std::vector<TLorentzVector>& v_clus, // TLV vector of all clusters of hard objects
                                      TLorentzVector& clus,                // TLV of current cluster
                                      TLorentzVector& HR,                  // uncorrected HR
                                      const float Drcone,                       // Cone size for el-pfo association
                                      const float MinDistCone,                  // Cone size for getting random Phi
                                      float& UEcorr) const                 // UE correction (result)
  {
      // 1. Get random phi
      unsigned int seed = floor( clus.Pt() * Gaudi::Units::GeV );
      TRandom3 hole;
      hole.SetSeed(seed);
 
      bool isNextToPart(true);
      bool isNextToHR(true);
      double phiRnd(0.);
 
      int numOfRndTrials = 0; // Counter for trials to find random cone without overlaps
      const int maxNumOfRndTrials = 100; // Max. number of trials to find random cone without overlaps
 
      while(isNextToPart || isNextToHR ){
        isNextToPart = false;
        isNextToHR = true;
 
        phiRnd = hole.Uniform( -std::numbers::pi, std::numbers::pi);
        double dR = P4Helpers::deltaR( HR.Eta(), HR.Phi(), clus.Eta(), phiRnd );
        if(dR > MinDistCone){
          isNextToHR = false;
        }
 
        for(const auto& clus_j : v_clus) { // loop over leptons
          dR = P4Helpers::deltaR( clus.Eta(), phiRnd, clus_j.Eta(), clus_j.Phi() );
          if(dR < MinDistCone){
            isNextToPart = true;
            break;
          }
        } // swclus_j
 
        numOfRndTrials++;
        if(numOfRndTrials == maxNumOfRndTrials){ // check number of trials
          UEcorr = 0.;
          return StatusCode::SUCCESS;
        }
      } // while isNextToPart, isNextToHR
 
      ATH_MSG_DEBUG("Found rnd phi: " << phiRnd);
 
 
      // 2. Calculete UE correction
      TLorentzVector tv_UEcorr; // TLV of UE correction (initialized with 0,0,0,0 automatically)
      std::pair <double, double> eta_rndphi = std::make_pair(clus.Eta(), phiRnd); // pair of current cluser eta and random phi
 
 
      // Calculate delta phi -> always the same angle so its sufficient to calculate it only once
      float dphi_angle=P4Helpers::deltaPhi(clus.Phi(),eta_rndphi.second);
 
      for(const auto fe_itr : *constits.feCont){ // loop over PFOs
        if(fe_itr->pt() < 0 || fe_itr->e() < 0){ //sanity check
          continue;
        }
 
        //remove charged FE that are not matched to the PV
        const static SG::ConstAccessor<char> PVMatchedAcc("matchedToPV");
        if(fe_itr->isCharged() && !PVMatchedAcc(*fe_itr)){
          continue;
        }
 
        double dR = P4Helpers::deltaR( fe_itr->eta(), fe_itr->phi(), eta_rndphi.first,  eta_rndphi.second );
        if( dR < Drcone ){
          // Rotate on dphi_angle
          TLorentzVector tv_fe = fe_itr->p4();
          tv_fe.RotateZ(dphi_angle);
          tv_UEcorr += tv_fe;  // summing PFOs of UE for correction
        } // cone requirement
      } // loop over PFOs
 
      UEcorr = tv_UEcorr.Pt();  // Pt of UE correction
 
      return StatusCode::SUCCESS;
  }
}