egammaTruthAssociationAlg.cxx

Go to the documentation of this file.

/*
  Copyright (C) 2002-2024 CERN for the benefit of the ATLAS collaboration
*/
 
#include "egammaTruthAssociationAlg.h"
#include "StoreGate/ReadHandle.h"
#include "StoreGate/WriteDecorHandle.h"
#include "StoreGate/WriteHandle.h"
#include "xAODCaloEvent/CaloClusterContainer.h"
#include "xAODEgamma/ElectronContainer.h"
#include "xAODEgamma/PhotonContainer.h"
#include "xAODEgamma/EgammaxAODHelpers.h"
#include "xAODTruth/xAODTruthHelpers.h"
 
#include "TruthUtils/HepMCHelpers.h"
 
#include <memory>
 
namespace {
xAOD::TruthParticle* getEgammaTruthParticle(
    const xAOD::TruthParticle* truth,
    xAOD::TruthParticleContainer& egammaTruthContainer) {
  if (!truth) {
    return nullptr;
  }
  // Find the original truth particle for electrons from conversions
  for (unsigned int i = 0;
       i < 100 && truth && HepMC::is_simulation_particle(truth); ++i) {
    if (truth->prodVtx() && truth->prodVtx()->nIncomingParticles()) {
      truth = truth->prodVtx()->incomingParticle(0);
    } else {
      break;
    }
  }
 
  // In case truth became null in the above loop
  if (!truth) {
    return nullptr;
  }
  for (auto egammaTruth : egammaTruthContainer) {
    if (HepMC::uniqueID(truth) == HepMC::uniqueID(*egammaTruth)) {
      return egammaTruth;
    }
  }
  return nullptr;
}
}  // namespace
 
using TruthLink_t = ElementLink<xAOD::TruthParticleContainer>;
using ClusterLink_t = ElementLink<xAOD::CaloClusterContainer>;
using ElectronLink_t = ElementLink<xAOD::ElectronContainer>;
using PhotonLink_t = ElementLink<xAOD::PhotonContainer>;
 
egammaTruthAssociationAlg::egammaTruthAssociationAlg(const std::string& name,
                                                     ISvcLocator* pSvcLocator)
  : AthReentrantAlgorithm(name, pSvcLocator)
{}
 
StatusCode
egammaTruthAssociationAlg::initialize()
{
 
  ATH_MSG_DEBUG("Initializing " << name() << "...");
 
  // initialize the data handles
  ATH_CHECK(m_truthEventContainerKey.initialize(m_doEgammaTruthContainer));
  ATH_CHECK(
    m_egammaTruthParticleContainerKey.initialize(m_doEgammaTruthContainer));
  ATH_CHECK(m_truthParticleContainerKey.initialize());
 
  // Now the standard decoration handles
  if (m_matchElectrons) {
    ATH_CHECK(initializeDecorKeys(m_electronDecKeys, m_electronDecName));
  } else {
    m_electronDecKeys.clear();
  }
 
  if (m_matchPhotons) {
    ATH_CHECK(initializeDecorKeys(m_photonDecKeys, m_photonDecName));
  } else {
    m_photonDecKeys.clear();
  }
 
  if (m_matchClusters) {
    ATH_CHECK(initializeDecorKeys(m_clusterDecKeys, m_clusterDecName));
  } else {
    m_clusterDecKeys.clear();
  }
 
  if (m_matchForwardElectrons) {
    ATH_CHECK(initializeDecorKeys(m_fwdElectronDecKeys, m_fwdElectronDecName));
  } else {
    m_fwdElectronDecKeys.clear();
  }
 
  CHECK(m_mcTruthClassifier.retrieve());
  ATH_MSG_DEBUG("Retrieved tool " << m_mcTruthClassifier);
  ATH_MSG_DEBUG("Initialization successful");
  return StatusCode::SUCCESS;
}
 
StatusCode
egammaTruthAssociationAlg::finalize()
{
  return StatusCode::SUCCESS;
}
 
StatusCode
egammaTruthAssociationAlg::execute(const EventContext& ctx) const
{
 
  SG::WriteHandle<xAOD::TruthParticleContainer> egammaTruthContainer;
 
  if (m_doEgammaTruthContainer) {
 
    egammaTruthContainer = SG::WriteHandle<xAOD::TruthParticleContainer>(
      m_egammaTruthParticleContainerKey, ctx);
    ATH_CHECK(egammaTruthContainer.record(
      std::make_unique<xAOD::TruthParticleContainer>(),
      std::make_unique<xAOD::TruthParticleAuxContainer>()));
 
    // Add a copy of electrons and photons to the truth egamma container
    SG::ReadHandle<xAOD::TruthEventContainer> truthEvtContainer(
      m_truthEventContainerKey, ctx);
 
    // only for serial running. Can remove check later
    if (!truthEvtContainer.isValid() || truthEvtContainer->empty()) {
      ATH_MSG_WARNING("Could not retrieve "
                      << m_truthEventContainerKey.key()
                      << " or container empty, returning");
      return StatusCode::SUCCESS;
    }
 
    for (const auto& truthParticleLink :
         truthEvtContainer->front()->truthParticleLinks()) {
      if (!truthParticleLink.isValid())
        continue;
      const xAOD::TruthParticle* truthParticle = *truthParticleLink;
      if (!isPromptEgammaParticle(ctx, truthParticle))
        continue;
      getNewTruthParticle(ctx,
                          *egammaTruthContainer,
                          truthParticle,
                          truthParticleLink.getDataPtr());
    }
  }
  // Decorate containers with truth info, including links to truth particles
  // Decorate the truth particles with links to the reco ones
 
  // note that in multithreading this must be valid; can't just fail with
  // success.
  SG::ReadHandle<xAOD::TruthParticleContainer> truthParticles(
    m_truthParticleContainerKey, ctx);
 
  // accessors
  static const SG::AuxElement::Accessor<ClusterLink_t> accClusLink(
    "recoClusterLink");
  static const SG::AuxElement::Accessor<ElectronLink_t> accElLink(
    "recoElectronLink");
  static const SG::AuxElement::Accessor<PhotonLink_t> accPhLink(
    "recoPhotonLink");
 
  if (m_matchElectrons) {
    ATH_MSG_DEBUG("About to match electrons");
    ATH_CHECK(match(ctx,
                    *truthParticles,
                    m_electronDecKeys,
                    accElLink,
                    egammaTruthContainer.ptr()));
  }
 
  if (m_matchPhotons) {
    ATH_MSG_DEBUG("About to match photons");
    ATH_CHECK(match(ctx,
                    *truthParticles,
                    m_photonDecKeys,
                    accPhLink,
                    egammaTruthContainer.ptr()));
  }
 
  if (m_matchClusters) {
    ATH_MSG_DEBUG("About to match clusters");
    ATH_CHECK(match(ctx,
                    *truthParticles,
                    m_clusterDecKeys,
                    accClusLink,
                    egammaTruthContainer.ptr()));
  }
 
  if (m_matchForwardElectrons) {
    ATH_MSG_DEBUG("About to match fwd electrons");
    ATH_CHECK(match(ctx,
                    *truthParticles,
                    m_fwdElectronDecKeys,
                    accElLink,
                    egammaTruthContainer.ptr()));
  }
 
  return StatusCode::SUCCESS;
}
 
bool
egammaTruthAssociationAlg::isPromptEgammaParticle(
  const EventContext& ctx,
  const xAOD::TruthParticle* truth) const
{
 
  if ((truth->pdgId() != 22 && abs(truth->pdgId()) != 11) ||
      MC::isDecayed(truth) || truth->status() == 3 ||
      HepMC::is_simulation_particle(truth) || truth->pt() < m_minPt) {
    return false;
  }
  MCTruthPartClassifier::Info mcinfo(ctx);
  auto type = m_mcTruthClassifier->particleTruthClassifier(truth, &mcinfo);
 
  // Isolated electron or photon
  if (type.first == MCTruthPartClassifier::IsoElectron ||
      type.first == MCTruthPartClassifier::IsoPhoton) {
    return true;
  }
 
  // FSR photon
  if (type.first == MCTruthPartClassifier::NonIsoPhoton &&
      type.second == MCTruthPartClassifier::FSRPhot &&
      truth->pt() > m_minPtFSR) {
    return true;
  }
 
  return false;
}
 
void
egammaTruthAssociationAlg::getNewTruthParticle(
  const EventContext& ctx,
  xAOD::TruthParticleContainer& egammaTruthContainer,
  const xAOD::TruthParticle* truth,
  const xAOD::TruthParticleContainer* oldContainer) const
{
  auto *truthParticle = egammaTruthContainer.push_back(std::make_unique<xAOD::TruthParticle>());
  truthParticle->setPdgId(truth->pdgId());
  truthParticle->setBarcode(HepMC::barcode(truth)); // FIXME barcode-based
  truthParticle->setStatus(truth->status());
  truthParticle->setPx(truth->px());
  truthParticle->setPy(truth->py());
  truthParticle->setPz(truth->pz());
  truthParticle->setE(truth->e());
  truthParticle->setM(truth->m());
  truthParticle->setProdVtxLink(truth->prodVtxLink());
  truthParticle->setDecayVtxLink(truth->decayVtxLink());
 
  static const SG::AuxElement::Accessor<ClusterLink_t> accClusLink(
    "recoClusterLink");
  static const SG::AuxElement::Accessor<ElectronLink_t> accElLink(
    "recoElectronLink");
  static const SG::AuxElement::Accessor<PhotonLink_t> accPhLink(
    "recoPhotonLink");
  static const SG::AuxElement::Accessor<TruthLink_t> accTruthLink(
    "truthParticleLink");
  static const SG::AuxElement::Accessor<int> accType("truthType");
  static const SG::AuxElement::Accessor<int> accOrigin("truthOrigin");
 
  if (m_matchClusters) {
    accClusLink(*truthParticle) = ClusterLink_t();
  }
  accElLink(*truthParticle) = ElectronLink_t();
  accPhLink(*truthParticle) = PhotonLink_t();
  accTruthLink(*truthParticle) = TruthLink_t(truth, *oldContainer, ctx);
  accTruthLink(*truthParticle).toPersistent();
  // MCTruthClassifier info
  MCTruthPartClassifier::Info mcinfo(ctx);
  auto info = m_mcTruthClassifier->particleTruthClassifier(truth, &mcinfo);
  accType(*truthParticle) = static_cast<int>(info.first);
  accOrigin(*truthParticle) = static_cast<int>(info.second);
}
 
template<class T>
StatusCode
egammaTruthAssociationAlg::initializeDecorKeys(
  SG::WriteDecorHandleKeyArray<T>& keys,
  const std::string& name)
{
  if (!keys.empty()) {
    ATH_MSG_FATAL("The WriteDecorHandle should not be configured directly.");
    return StatusCode::FAILURE;
  }
 
  keys.emplace_back(name + ".truthParticleLink");
  keys.emplace_back(name + ".truthType");
  keys.emplace_back(name + ".truthOrigin");
  ATH_CHECK(keys.initialize());
  return StatusCode::SUCCESS;
}
 
// constructor
template<class T>
egammaTruthAssociationAlg::writeDecorHandles<T>::writeDecorHandles(
  const SG::WriteDecorHandleKeyArray<T>& hkeys,
  const EventContext& ctx)
  : el(hkeys.at(0), ctx)
  , type(hkeys.at(1), ctx)
  , origin(hkeys.at(2), ctx)
{}
 
template<class T>
egammaTruthAssociationAlg::MCTruthInfo_t
egammaTruthAssociationAlg::particleTruthClassifier(
  const EventContext& ctx,
  const T* particle) const
{
  MCTruthInfo_t info{};
  MCTruthPartClassifier::Info mcinfo(ctx);
  auto ret = m_mcTruthClassifier->particleTruthClassifier(particle, &mcinfo);
  info.genPart = mcinfo.genPart;
  info.first = ret.first;
  info.second = ret.second;
  return info;
}
 
template<>
egammaTruthAssociationAlg::MCTruthInfo_t
egammaTruthAssociationAlg::particleTruthClassifier<xAOD::Electron>(
  const EventContext& ctx,
  const xAOD::Electron* electron) const
{
  MCTruthInfo_t info{};
  MCTruthPartClassifier::Info mcinfo(ctx);
  auto ret = m_mcTruthClassifier->particleTruthClassifier(electron, &mcinfo);
  if (ret.first == MCTruthPartClassifier::Unknown &&
      !xAOD::EgammaHelpers::isFwdElectron(electron) &&
      electron->caloCluster()) {
    ATH_MSG_DEBUG("Trying cluster-based truth classification for electron");
    ret = m_mcTruthClassifier->particleTruthClassifier(electron->caloCluster(),
                                                       &mcinfo);
  }
  info.genPart = mcinfo.genPart;
  info.first = ret.first;
  info.second = ret.second;
  return info;
}
 
template<class T, class L>
StatusCode
egammaTruthAssociationAlg::match(
  const EventContext& ctx,
  const xAOD::TruthParticleContainer& truthParticles,
  const SG::WriteDecorHandleKeyArray<T>& hkeys,
  const SG::AuxElement::Accessor<L>& linkAccess,
  xAOD::TruthParticleContainer* egammaTruthContainer) const
{
 
  writeDecorHandles<T> decoHandles(hkeys, ctx);
 
 
  for (auto particle : *decoHandles.readHandle()) {
 
    MCTruthInfo_t info =
      particleTruthClassifier(ctx, particle);
 
    const xAOD::TruthParticle* truthParticle = info.genPart;
    if (truthParticle) {
      ElementLink<xAOD::TruthParticleContainer> link(
        truthParticle, truthParticles, ctx);
      ATH_MSG_DEBUG(
        "Decorating object with link to truth, index = " << link.index());
      decoHandles.el(*particle) = link;
    } else {
      decoHandles.el(*particle) = ElementLink<xAOD::TruthParticleContainer>();
    }
    decoHandles.el(*particle).toPersistent();
    ATH_MSG_DEBUG("truthType = " << info.first
                                 << " truthOrigin = " << info.second);
    decoHandles.type(*particle) = static_cast<int>(info.first);
    decoHandles.origin(*particle) = static_cast<int>(info.second);
 
    // Decorate the corresponding truth particle with the link to the reco
    if (m_doEgammaTruthContainer) {
      if (!egammaTruthContainer) {
        ATH_MSG_ERROR("The egammaTruthContainer needs to be valid");
        return StatusCode::FAILURE;
      }
      const xAOD::TruthParticle* truth =
        xAOD::TruthHelpers::getTruthParticle(*particle);
      if (truth) {
        xAOD::TruthParticle* truthEgamma =
          getEgammaTruthParticle(truth, *egammaTruthContainer);
        if (truthEgamma) {
          // we found a truthEgamma object we should annotate if this is the
          // best link
          bool annotateLink = true;                   // by default we annotate
          const auto link = linkAccess(*truthEgamma); // what already exists
          if (link.isValid()) {
            auto oldPart = *link;
            if (oldPart && truthEgamma->e() > 0 &&
                std::abs(oldPart->e() / truthEgamma->e() - 1) <
                  std::abs(particle->e() / truthEgamma->e() - 1)) {
              ATH_MSG_DEBUG(truthEgamma
                            << ": "
                            << " already set to a better matched particle: "
                            << particle);
              annotateLink = false;
            }
          }
 
          if (annotateLink) {
            L link(particle, *decoHandles.readHandle(), ctx);
            linkAccess(*truthEgamma) = link;
            linkAccess(*truthEgamma).toPersistent();
          }
        }
      }
    }
  }
  return StatusCode::SUCCESS;
}