VertexDecoratorAlg.cxx

Go to the documentation of this file.

/*
  Copyright (C) 2002-2025 CERN for the benefit of the ATLAS collaboration
*/
 
#include "VertexDecoratorAlg.h"
#include "xAODEgamma/ElectronxAODHelpers.h"
#include "AthContainers/AuxElement.h"
#include "AsgDataHandles/WriteDecorHandle.h"
#include "AsgDataHandles/ReadDecorHandle.h"
#include "PhotonVertexSelection/PhotonVertexHelpers.h"
#include <StoreGate/WriteHandleKey.h>
#include <xAODParticleEvent/CompositeParticleAuxContainer.h>
#include <TLorentzVector.h>
#include <cmath>
 
 
namespace InDetGNNHardScatterSelection
{
  VertexDecoratorAlg ::VertexDecoratorAlg(const std::string &name,
                                          ISvcLocator *pSvcLocator)
      : AthHistogramAlgorithm (name, pSvcLocator)
  {
  }
 
  StatusCode VertexDecoratorAlg ::initialize()
  {
    ATH_CHECK(m_vertexInKey.initialize());
    ATH_CHECK(m_eventInKey.initialize());
    ATH_CHECK(m_photonsInKey.initialize());
    ATH_CHECK(m_electronsInKey.initialize());
    ATH_CHECK(m_muonsInKey.initialize());
    ATH_CHECK(m_jetsInKey.initialize());
    ATH_CHECK(m_multiPhotonsOutKey.initialize());
 
    const std::string baseName = m_vertexInKey.key();
 
    // WriteHandleKeys
    m_photonLinksKey = baseName + ".photonLinks";
    m_jetLinksKey = baseName + ".jetLinks";
    m_electronLinksKey = baseName + ".electronLinks";
    m_muonLinksKey = baseName + ".muonLinks";
    m_mDecor_ntrk = baseName + "." + m_mDecor_ntrk.key();
    m_mDecor_sumPt = baseName + "." + m_mDecor_sumPt.key();
    m_mDecor_chi2Over_ndf = baseName + "." + m_mDecor_chi2Over_ndf.key();
    m_mDecor_z_asym = baseName + "." + m_mDecor_z_asym.key();
    m_mDecor_weighted_z_asym = baseName + "." + m_mDecor_weighted_z_asym.key();
    m_mDecor_weighted_z_kurt = baseName + "." + m_mDecor_weighted_z_kurt.key();
    m_mDecor_z_skew = baseName + "." + m_mDecor_z_skew.key();
    m_mDecor_photon_deltaz = baseName + "." + m_mDecor_photon_deltaz.key();
    m_mDecor_photon_deltaPhi = baseName + "." + m_mDecor_photon_deltaPhi.key();
    m_mDecor_actualInterPerXing = baseName + "." + m_mDecor_actualInterPerXing.key();
    m_multiPhotonLinksKey = baseName + ".multiPhotonLinks";
    m_mDecor_gnnScore = baseName + "." + m_mDecor_gnnScore.key();
    m_mDecor_nElectrons = baseName + "." + m_mDecor_nElectrons.key();
    m_mDecor_nMuons = baseName + "." + m_mDecor_nMuons.key();
    m_mDecor_nPhotons = baseName + "." + m_mDecor_nPhotons.key();
    m_mDecor_nJets = baseName + "." + m_mDecor_nJets.key();
 
    ATH_CHECK(m_multiPhotonLinksKey.initialize());
    ATH_CHECK(m_photonLinksKey.initialize());
    ATH_CHECK(m_jetLinksKey.initialize());
    ATH_CHECK(m_electronLinksKey.initialize());
    ATH_CHECK(m_muonLinksKey.initialize());
    ATH_CHECK(m_mDecor_ntrk.initialize());
    ATH_CHECK(m_mDecor_sumPt.initialize());
    ATH_CHECK(m_mDecor_chi2Over_ndf.initialize());
    ATH_CHECK(m_mDecor_z_asym.initialize());
    ATH_CHECK(m_mDecor_weighted_z_asym.initialize());
    ATH_CHECK(m_mDecor_weighted_z_kurt.initialize());
    ATH_CHECK(m_mDecor_z_skew.initialize());
    ATH_CHECK(m_mDecor_photon_deltaz.initialize());
    ATH_CHECK(m_mDecor_photon_deltaPhi.initialize());
    ATH_CHECK(m_mDecor_actualInterPerXing.initialize());
    ATH_CHECK(m_mDecor_gnnScore.initialize());
    ATH_CHECK(m_mDecor_nElectrons.initialize());
    ATH_CHECK(m_mDecor_nMuons.initialize());
    ATH_CHECK(m_mDecor_nJets.initialize());
    ATH_CHECK(m_mDecor_nPhotons.initialize());
 
    // ReadHandleKeys
    m_deltaZKey = baseName + ".deltaZ";
    m_deltaPhiKey = baseName + ".deltaPhi";
    ATH_CHECK(m_deltaZKey.initialize());
    ATH_CHECK(m_deltaPhiKey.initialize());
 
    // additional ReadHandleKeys to declare dependencies to the scheduler
    std::string photonBaseName = m_photonsInKey.key();
 
    m_caloPointingZKey = photonBaseName + ".caloPointingZ";
    m_zCommonKey = photonBaseName + ".zCommon";
    m_zCommonErrorKey = photonBaseName + ".zCommonError";
 
    ATH_CHECK(m_caloPointingZKey.initialize());
    ATH_CHECK(m_zCommonKey.initialize());
    ATH_CHECK(m_zCommonErrorKey.initialize());
 
    // Tools
    ATH_CHECK(m_gnnTool.retrieve());
    ATH_CHECK(m_trkVtxAssociationTool.retrieve());
 
    return StatusCode::SUCCESS;
  }
 
  StatusCode VertexDecoratorAlg ::execute(const EventContext& ctx) 
  {
    SG::ReadHandle<xAOD::VertexContainer> vertices(m_vertexInKey, ctx);
    ATH_CHECK(vertices.isValid());
 
    SG::ReadHandle<xAOD::PhotonContainer> photonsIn(m_photonsInKey, ctx);
    ATH_CHECK(photonsIn.isValid());
    
    SG::ReadHandle<xAOD::ElectronContainer> electronsIn(m_electronsInKey, ctx);
    ATH_CHECK(electronsIn.isValid());
    SG::ReadHandle<xAOD::MuonContainer> muonsIn(m_muonsInKey, ctx);
    ATH_CHECK(muonsIn.isValid());
    SG::ReadHandle<xAOD::JetContainer> jetsIn(m_jetsInKey, ctx);
    ATH_CHECK(jetsIn.isValid());
    SG::ReadHandle<xAOD::EventInfo> eventInfo(m_eventInKey,ctx);
    ATH_CHECK(eventInfo.isValid());
 
    SG::WriteDecorHandle<xAOD::VertexContainer, std::vector<ElementLink<xAOD::PhotonContainer>>> dec_photonLinks(m_photonLinksKey, ctx);
    SG::WriteDecorHandle<xAOD::VertexContainer, std::vector<ElementLink<xAOD::JetContainer>>> dec_jetLinks(m_jetLinksKey, ctx);
    SG::WriteDecorHandle<xAOD::VertexContainer, std::vector<ElementLink<xAOD::ElectronContainer>>> dec_electronLinks(m_electronLinksKey, ctx);
    SG::WriteDecorHandle<xAOD::VertexContainer, std::vector<ElementLink<xAOD::MuonContainer>>> dec_muonLinks(m_muonLinksKey, ctx);
    SG::WriteDecorHandle<xAOD::VertexContainer, std::vector<ElementLink<xAOD::CompositeParticleContainer>>> dec_multiPhotonLinks(m_multiPhotonLinksKey, ctx);
 
    SG::ReadDecorHandle<xAOD::VertexContainer, float> acc_deltaZ(m_deltaZKey, ctx);
    SG::ReadDecorHandle<xAOD::VertexContainer, float> acc_deltaPhi(m_deltaPhiKey, ctx);
    SG::ReadDecorHandle<xAOD::PhotonContainer, float> acc_caloPointingZ(m_caloPointingZKey, ctx);
    SG::ReadDecorHandle<xAOD::PhotonContainer, float> acc_zCommon(m_zCommonKey, ctx);
    SG::ReadDecorHandle<xAOD::PhotonContainer, float> acc_zCommonError(m_zCommonErrorKey, ctx);
 
    SG::WriteDecorHandle<xAOD::VertexContainer, float> dec_gnnScore(m_mDecor_gnnScore,ctx);
    
    // Decorations needed by the GNNTool
    SG::WriteDecorHandle<xAOD::VertexContainer, int> dec_ntrk(m_mDecor_ntrk, ctx);
    SG::WriteDecorHandle<xAOD::VertexContainer, float> dec_sumPt(m_mDecor_sumPt,ctx);
    SG::WriteDecorHandle<xAOD::VertexContainer, float> dec_chi2Over_ndf(m_mDecor_chi2Over_ndf,ctx);
    SG::WriteDecorHandle<xAOD::VertexContainer, float> dec_z_asym(m_mDecor_z_asym,ctx);
    SG::WriteDecorHandle<xAOD::VertexContainer, float> dec_weighted_z_asym(m_mDecor_weighted_z_asym,ctx);
    SG::WriteDecorHandle<xAOD::VertexContainer, float> dec_z_kurt(m_mDecor_weighted_z_kurt,ctx);
    SG::WriteDecorHandle<xAOD::VertexContainer, float> dec_z_skew(m_mDecor_z_skew,ctx);
    SG::WriteDecorHandle<xAOD::VertexContainer, float> dec_photon_deltaz(m_mDecor_photon_deltaz,ctx);
    SG::WriteDecorHandle<xAOD::VertexContainer, float> dec_photon_deltaPhi(m_mDecor_photon_deltaPhi,ctx);
    SG::WriteDecorHandle<xAOD::VertexContainer, float> dec_actualInterPerXing(m_mDecor_actualInterPerXing,ctx);
    SG::WriteDecorHandle<xAOD::VertexContainer, int> dec_nElectrons(m_mDecor_nElectrons,ctx);
    SG::WriteDecorHandle<xAOD::VertexContainer, int> dec_nMuons(m_mDecor_nMuons,ctx);
    SG::WriteDecorHandle<xAOD::VertexContainer, int> dec_nJets(m_mDecor_nJets,ctx);
    SG::WriteDecorHandle<xAOD::VertexContainer, int> dec_nPhotons(m_mDecor_nPhotons,ctx);
    SG::WriteHandle<xAOD::CompositeParticleContainer> mpHandle(m_multiPhotonsOutKey, ctx);
 
    auto mpCont = std::make_unique<xAOD::CompositeParticleContainer>();
    auto mpAux  = std::make_unique<xAOD::CompositeParticleAuxContainer>();
    mpCont->setStore(mpAux.get());
 
    const xAOD::Vertex* bestVtxForMP = nullptr;
    ElementLink<xAOD::CompositeParticleContainer> mpLink; // will be valid only if created
    bool haveMP = false;
 
    // For MT safety, decoration keys must be initialized and available upfront.
    const bool haveZCommonDecor = acc_zCommon.isAvailable() && acc_zCommonError.isAvailable();
    const bool haveCaloPointingDecor = acc_caloPointingZ.isAvailable();
    
    if (photonsIn->size() > 1 && (!haveZCommonDecor || !haveCaloPointingDecor))
    {
      ATH_MSG_ERROR("photonsIn has size > 1 but required decorations are missing: "
                    << m_zCommonKey.key() << " and/or " << m_zCommonErrorKey.key()
                    << " and/or " << m_caloPointingZKey.key());
      return StatusCode::FAILURE;
    }
 
    if (photonsIn->size() > 1)
    {
      // Combined pointing Z (weighted)
      double sumW = 0.0;
      double sumWZ = 0.0;
 
      TLorentzVector p4sum;
      int nUsed = 0;
      for (const xAOD::Photon* ph : *photonsIn)
      {
        const float zCommon = acc_zCommon(*ph);
        const float zCaloPointing = acc_caloPointingZ(*ph);
        // Value-based selection only: prefer zCommon, then fallback to caloPointing value.
        const bool useZCommonValue = std::isfinite(zCommon);
        const float z = useZCommonValue ? zCommon : zCaloPointing;
        if (!std::isfinite(z)) {
          ATH_MSG_ERROR("Non-finite photon pointing values found in both "
                        << m_zCommonKey.key() << " and " << m_caloPointingZKey.key());
          return StatusCode::FAILURE;
        }
 
        // Keep weighting consistent with the z source:
        // use zCommonError only when z comes from zCommon.
        double w = 1.0;
        if (useZCommonValue) {
          const float s = acc_zCommonError(*ph);
          if (std::isfinite(s) && s > 0.0f) w = 1.0 / (double(s) * double(s));
        }
        
        sumW  += w;
        sumWZ += w * z;
        p4sum += ph->p4();
        ++nUsed;
      }
 
      if (nUsed > 1 && sumW > 0.0)
      {
        const float zPoint = sumWZ / sumW;
 
        // Choose vertex closest in z to combined pointing
        float bestAbsDZ = 1e30;
        for (const xAOD::Vertex* vtx : *vertices)
        {
          if (vtx->vertexType() == xAOD::VxType::NoVtx) continue;
 
          const float dz = zPoint - vtx->z();
          const float adz = std::abs(dz);
          if (adz < bestAbsDZ)
          {
            bestAbsDZ = adz;
            bestVtxForMP = vtx;
          }
        }
 
        if (bestVtxForMP)
        {
          // Create exactly one CompositeParticle node
          auto* mp = new xAOD::CompositeParticle();
          mpCont->push_back(mp);
 
          mp->setP4(p4sum);
 
          SG::AuxElement::Decorator<float> dec_mp_deltaZ("deltaZ");
          SG::AuxElement::Decorator<float> dec_mp_deltaPhi("deltaPhi");
          SG::AuxElement::Decorator<int>   dec_mp_nPhotons("nPhotons");
          SG::AuxElement::Decorator<float> dec_mp_zPointing("zPointing");
 
          const float deltaZ = zPoint - bestVtxForMP->z();
          dec_mp_deltaZ(*mp) = deltaZ;
          dec_mp_zPointing(*mp) = zPoint;
          dec_mp_nPhotons(*mp) = static_cast<int>(photonsIn->size());
 
          float dphi = 0.0f;
          if (acc_deltaPhi.isAvailable())
          {
            dphi = acc_deltaPhi(*bestVtxForMP);
            if (!std::isfinite(dphi)) dphi = 0.0f;
          }
          dec_mp_deltaPhi(*mp) = dphi;
          mpLink.setElement(mp);
          mpLink.setStorableObject(*mpCont, true);
 
          haveMP = true;
        }
      }
    }
 
    // Record the MultiPhotons container (even if empty)
    ATH_CHECK(mpHandle.record(std::move(mpCont), std::move(mpAux)));
 
    std::map< const xAOD::Vertex*, std::vector<ElementLink<xAOD::JetContainer>> > jetsInVertex;
    std::map< const xAOD::Jet*, std::map< const xAOD::Vertex*, int> > jetVertexPt;
 
    // initialize jet-vertex map 
    for(const xAOD::Vertex *vertex : *vertices){
      jetsInVertex[vertex] = {};
      for(const xAOD::Jet *jet : *jetsIn){
        jetVertexPt[jet][vertex] = 0;
      }
    }
 
    // pre-fill the jet-vertex map
    for(const xAOD::Jet* jet : *jetsIn){
 
      // for each jet, calculate the track pT associated to each vertex
      std::vector<const xAOD::TrackParticle*> ghostTracks = jet->getAssociatedObjects<xAOD::TrackParticle >(xAOD::JetAttribute::GhostTrack);
      for(const xAOD::TrackParticle* jtrk : ghostTracks){
        if( !jtrk ) continue;
        auto jetTrackVertex = m_trkVtxAssociationTool->getUniqueMatchVertexLink(*jtrk, *vertices);
        if(jetTrackVertex) jetVertexPt[jet][*jetTrackVertex] += jtrk->pt();
      }
 
      // find vertex with the largest fraction of jet track pt
      float maxPtFrac = -1;
      const xAOD::Vertex* uniqueVertexAddress = nullptr;
      for (const xAOD::Vertex *vertex : *vertices) {
        if (vertex->vertexType() == xAOD::VxType::NoVtx) continue;
        if(jetVertexPt[jet][vertex] > maxPtFrac){
          maxPtFrac = jetVertexPt[jet][vertex];
          uniqueVertexAddress = vertex;
        }
      }
 
      // add jet to that vertex's vector of links
      ElementLink<xAOD::JetContainer> jetLink;
      jetLink.setElement(jet);
      jetLink.setStorableObject(*jetsIn.ptr(), true);
      jetsInVertex[uniqueVertexAddress].push_back(jetLink);
    }
 
    for (const xAOD::Vertex *vertex : *vertices)
    {
      if (vertex->vertexType() == xAOD::VxType::NoVtx)
        continue;
 
      if (vertex->nTrackParticles() < 2)
        continue;
 
      dec_actualInterPerXing(*vertex) = eventInfo->actualInteractionsPerCrossing();
 
      // variables for calculation of delta Z asymmetry and delta d asymmetry
      float z_asym = 0;
      float sumDZ = 0;
      float deltaZ = 0;
      float modsumDZ = 0;
      float weighted_sumDZ = 0;
      float weighted_deltaZ = 0;
      float weighted_modsumDZ = 0;
      float weighted_z_asym = 0;
 
      // make vector
      std::vector<float> track_deltaZ;
 
      for (size_t i = 0; i < vertex->nTrackParticles(); i++) {
 
        const xAOD::TrackParticle *trackTmp = vertex->trackParticle(i);
 
        if(!trackTmp) continue;
 
        deltaZ = trackTmp->z0() + trackTmp->vz() - vertex->z();
        track_deltaZ.push_back(deltaZ);
        // get the track weight for each track to get the deltaZ/trk_weight
        float trk_weight = vertex->trackWeight(i);
        weighted_deltaZ = deltaZ * trk_weight;
        // sum of delta z
        sumDZ += deltaZ;
        modsumDZ += std::abs(deltaZ);
        weighted_sumDZ += weighted_deltaZ;
        weighted_modsumDZ += std::abs(weighted_deltaZ);
      } // end loop over tracks
 
      if (modsumDZ > 0) {
        z_asym = sumDZ / modsumDZ;
      }
      if (weighted_modsumDZ > 0) {
        weighted_z_asym = weighted_sumDZ / weighted_modsumDZ;
      }
 
      const float number_tracks = track_deltaZ.size(); // get number of tracks
      const float mean_Dz =
        number_tracks > 0 ? sumDZ / number_tracks : 0.F; // calculate average
 
      float z_skew = 0; // skewness of DeltaZ asymmetry
      float z_kurt = 0; // Kurtosis of DeltaZ asymmetry
      float z_var = 0;  // variance of DeltaZ
 
      for (auto i : track_deltaZ)
      {
        float z_zbar = (i - mean_Dz);
        z_var += std::pow(z_zbar, 2);
        z_skew += std::pow(z_zbar, 3);
        z_kurt += std::pow(z_zbar, 4);
      }
      if (number_tracks > 1 && z_var > 0) {
        z_var /= (number_tracks - 1);
        float z_sd = std::sqrt(z_var);
        const float skew_denom = (number_tracks - 1) * std::pow(z_sd, 3);
        const float kurt_denom = (number_tracks - 1) * std::pow(z_sd, 4);
        if (std::isfinite(skew_denom) && skew_denom != 0.F) {
          z_skew /= skew_denom;
        } else {
          z_skew = 0.F;
        }
        if (std::isfinite(kurt_denom) && kurt_denom != 0.F) {
          z_kurt /= kurt_denom;
        } else {
          z_kurt = 0.F;
        }
      }
      else
      {
        ATH_MSG_WARNING("z momenta are NaN: setting to zero");
        z_skew = 0.;
        z_kurt = 0.;
      }
 
      dec_ntrk(*vertex) = number_tracks;
 
      if(!dec_sumPt.isAvailable()){
        dec_sumPt(*vertex) = xAOD::PVHelpers::getVertexSumPt(vertex, 1, false);
      }
      const float numberDoF = vertex->numberDoF();
      dec_chi2Over_ndf(*vertex) =
        numberDoF > 0.F ? vertex->chiSquared() / numberDoF : 0.F;
      dec_z_asym(*vertex) = z_asym;
      dec_weighted_z_asym(*vertex) = weighted_z_asym;
      dec_z_kurt(*vertex) = z_kurt;
      dec_z_skew(*vertex) = z_skew;
 
      if (acc_deltaZ.isAvailable() && photonsIn->size() > 0 ) {
        //protect against rare NaNs before assigning decorator: setting to 0 (-999 cause NaNs)
        if (!std::isfinite(acc_deltaZ(*vertex))) {
          ATH_MSG_WARNING("photon deltaPhi is NaN: setting to -1!");
          dec_photon_deltaz(*vertex) = -1;
        }
        else{
        dec_photon_deltaz(*vertex) = acc_deltaZ(*vertex);
        }
      }
      else{
       dec_photon_deltaz(*vertex) = -1;
      }
      if (acc_deltaPhi.isAvailable() && photonsIn->size() > 0) {
        
        if (!std::isfinite(acc_deltaPhi(*vertex))) {
          ATH_MSG_WARNING("photon deltaPhi is NaN: setting to 0!");
          dec_photon_deltaPhi(*vertex) = -1;
        }
        else{
        dec_photon_deltaPhi(*vertex) = acc_deltaPhi(*vertex);
        }
      }
      else{
       dec_photon_deltaPhi(*vertex) = -1;
      }
 
      // associate objects to vertices
      auto countValid = [](const auto& links) -> int {
        int n = 0;
        for (const auto& l : links) if (l.isValid()) ++n;
        return n;
      };
 
      std::vector<ElementLink<xAOD::ElectronContainer>> electronLinks;
      for(const xAOD::Electron* electron : *electronsIn){
        const auto *id_trk = xAOD::EgammaHelpers::getOriginalTrackParticle(electron);
        if(!id_trk) continue;
        auto eleVertex = m_trkVtxAssociationTool->getUniqueMatchVertexLink(*id_trk, *vertices);
        if(!eleVertex) continue;
        ElementLink<xAOD::ElectronContainer> elLink;
        if(*eleVertex == vertex){
            elLink.setElement(electron);
            elLink.setStorableObject(*electronsIn.ptr(), true);
            electronLinks.push_back(elLink);
        }
      }
      dec_electronLinks(*vertex) = electronLinks;
      dec_nElectrons(*vertex) = countValid(electronLinks);
 
      std::vector<ElementLink<xAOD::PhotonContainer>> photonLinks;
      for(const xAOD::Photon* photon : *photonsIn){
        ElementLink<xAOD::PhotonContainer> phLink;
        phLink.setElement(photon);
        phLink.setStorableObject(*photonsIn.ptr(), true);
        photonLinks.push_back(phLink);
      }
      dec_photonLinks(*vertex) = photonLinks;
      dec_nPhotons(*vertex) = countValid(photonLinks);
      
      // multi-photon link 
      std::vector<ElementLink<xAOD::CompositeParticleContainer>> mpLinks;
      if (haveMP && vertex == bestVtxForMP)
      {
        mpLinks.push_back(mpLink);
      }
      dec_multiPhotonLinks(*vertex) = mpLinks;
 
      // for jets, use prefilled map
      dec_jetLinks(*vertex) = jetsInVertex[vertex];
      dec_nJets(*vertex) = countValid(jetsInVertex[vertex]);
 
      std::vector<ElementLink<xAOD::MuonContainer>> muonLinks;
      for(const xAOD::Muon* muon : *muonsIn){
        const auto *tp = muon->trackParticle(xAOD::Muon::InnerDetectorTrackParticle);
        if(!tp) continue;
        try{
          auto muonVertex = m_trkVtxAssociationTool->getUniqueMatchVertexLink(*tp, *vertices);
          if(!muonVertex) continue;
          ElementLink<xAOD::MuonContainer> muonLink;
          if(*muonVertex == vertex){
            muonLink.setElement(muon);
            muonLink.setStorableObject(*muonsIn.ptr(), true);
            muonLinks.push_back(muonLink);
          }
        }catch(...) {
          ATH_MSG_DEBUG("Skipping muon as the track is not associated to any PV ");
          ATH_MSG_DEBUG("Muon pT, eta = " << muon->pt() << " " << muon->eta());
        }
 
      }
      dec_muonLinks(*vertex) = muonLinks;
      dec_nMuons(*vertex) = countValid(muonLinks);
 
      // Finally, decorate the vertices with the GNN score
      float score_phsvertex = m_gnnTool->decorate(*vertex);
 
      dec_gnnScore(*vertex) = score_phsvertex;
    }
 
    return StatusCode::SUCCESS;
  }
 
} // namespace InDetGNNHardScatterSelection