InDetPhysValTruthDecoratorAlg.cxx

Go to the documentation of this file.

/*
  Copyright (C) 2002-2024 CERN for the benefit of the ATLAS collaboration
*/

 
#include "InDetPhysValTruthDecoratorAlg.h"
#include "safeDecorator.h"
#include <limits>
 
#include "xAODTruth/TruthVertex.h"
#include "StoreGate/ReadDecorHandle.h"
#include "GaudiKernel/PhysicalConstants.h"
 
#include "TDatabasePDG.h"
#include "TParticlePDG.h"
#include "TrkParameters/TrackParameters.h" // Contains typedef to Trk::CurvilinearParameters
#include <cmath>
 
// ref:
// https://svnweb.cern.ch/trac/atlasoff/browser/Tracking/TrkEvent/TrkParametersBase/trunk/TrkParametersBase/CurvilinearParametersT.h
 
InDetPhysValTruthDecoratorAlg::InDetPhysValTruthDecoratorAlg(const std::string& name, ISvcLocator* pSvcLocator) :
  AthReentrantAlgorithm(name, pSvcLocator)
{
}
 
InDetPhysValTruthDecoratorAlg::~InDetPhysValTruthDecoratorAlg () {
  // nop
}
 
StatusCode
InDetPhysValTruthDecoratorAlg::initialize() {
  ATH_CHECK(m_extrapolator.retrieve());
  ATH_CHECK(m_beamSpotDecoKey.initialize());
  ATH_CHECK( m_truthPixelClusterName.initialize() );
  ATH_CHECK( m_truthSCTClusterName.initialize() );
  ATH_CHECK( m_truthSelectionTool.retrieve( EnableTool { not m_truthSelectionTool.name().empty() } ) );
  if (not m_truthSelectionTool.name().empty() ) {
    m_cutFlow = CutFlow(m_truthSelectionTool->nCuts() );
  }
 
  ATH_CHECK( m_truthParticleName.initialize());
  ATH_CHECK( m_truthParticleIndexDecor.initialize( !m_truthParticleIndexDecor.key().empty()));
 
  ATH_CHECK(m_truthEventName.initialize(!m_truthEventName.key().empty()));
  ATH_CHECK(m_truthPileupEventName.initialize(!m_truthPileupEventName.key().empty()));
 
  std::vector<std::string> decor_names(kNDecorators);
  decor_names[kDecorD0]="d0";
  decor_names[kDecorZ0]="z0";
  decor_names[kDecorPhi]="phi";
  decor_names[kDecorTheta]="theta";
  decor_names[kDecorZ0st]="z0st";
  decor_names[kDecorQOverP]="qOverP";
  decor_names[kDecorProdR]="prodR";
  decor_names[kDecorProdZ]="prodZ";
  decor_names[kDecorNSilHits]="nSilHits";
  decor_names[kDecorTime]="time";
 
  IDPVM::createDecoratorKeysAndAccessor(*this, m_truthParticleName,m_prefix.value(),decor_names, m_decor);
  assert( m_decor.size() == kNDecorators);
  return StatusCode::SUCCESS;
}
 
StatusCode
InDetPhysValTruthDecoratorAlg::finalize() {
  if (not m_truthSelectionTool.name().empty()) {
    std::lock_guard<std::mutex> lock(m_mutex);
    ATH_MSG_DEBUG( "Truth selection cut flow : " << m_cutFlow.report(m_truthSelectionTool->names()) );
  }
  if (m_nMissingTruthParticles>0) {
    ATH_MSG_INFO( "Clusters which reference missing / thinned truth particles : " << m_nMissingTruthParticles );
  }
  return StatusCode::SUCCESS;
}
 
StatusCode
InDetPhysValTruthDecoratorAlg::execute(const EventContext &ctx) const {
  SG::ReadHandle<xAOD::TruthParticleContainer> ptruth(m_truthParticleName, ctx);
  if ((not ptruth.isValid())) {
    return StatusCode::FAILURE;
  }
 
  std::size_t ptruth_size=ptruth->size();
 
  std::vector<unsigned int> truthIndexMap;
  if (!m_truthParticleIndexDecor.empty()) {
     SG::ReadDecorHandle<xAOD::TruthParticleContainer, unsigned int> decor_index(m_truthParticleIndexDecor, ctx);
     unsigned int max_size=0;
     assert( ptruth_size < std::numeric_limits<unsigned int>::max());
     for (const xAOD::TruthParticle *truth_particle : *ptruth) {
        max_size = std::max( max_size, decor_index(*truth_particle) );
     }
     if (max_size>=std::numeric_limits<unsigned int>::max()) {
        ATH_MSG_ERROR("Truth index exceed max allowed range.");
        return StatusCode::FAILURE;
     }
     ++max_size;
     truthIndexMap.resize( max_size, std::numeric_limits<unsigned int>::max());
     unsigned int new_index=0;
     for (const xAOD::TruthParticle *truth_particle : *ptruth) {
        truthIndexMap.at( decor_index(*truth_particle) ) = new_index;
        ++new_index;
     }
  }
  else {
     truthIndexMap.reserve(ptruth_size);
     for (unsigned int i=0; i<ptruth_size; ++i) {
        truthIndexMap.push_back(i);
     }
  }
 
  std::vector< IDPVM::OptionalDecoration<xAOD::TruthParticleContainer,float> >
     float_decor( IDPVM::createDecoratorsIfNeeded(*ptruth, m_decor, ctx, msgLvl(MSG::DEBUG)) );

  std::vector< std::array<uint16_t, kNClusterTypes> > tp_clustercount;
  tp_clustercount.resize(ptruth_size,std::array<uint16_t,kNClusterTypes>{});
  unsigned int missing_truth_particle=0u;
  //Loop over the pixel and sct clusters to fill the truth barcode - cluster count maps
  SG::ReadHandle<xAOD::TrackMeasurementValidationContainer> sctClusters(m_truthSCTClusterName, ctx); 
  SG::ReadHandle<xAOD::TrackMeasurementValidationContainer> pixelClusters(m_truthPixelClusterName, ctx); 
  //only decorate the truth particles with truth silicon hits if both containers are available 
  if (sctClusters.isValid() && pixelClusters.isValid()) {
    for (const auto *const sct : *sctClusters) {
      const xAOD::TrackMeasurementValidation* sctCluster = sct;
      static const SG::AuxElement::ConstAccessor< std::vector<unsigned int> > truthIndexAcc("truth_index");
      if (truthIndexAcc.isAvailable(*sctCluster)) {
        const std::vector<unsigned int> &truth_indices = truthIndexAcc(*sctCluster);
        for (auto index : truth_indices) {
           if (index != std::numeric_limits<unsigned int>::max()) {
              if (index < truthIndexMap.size() && truthIndexMap[index] != std::numeric_limits<unsigned int>::max()) {
                 ++tp_clustercount.at(truthIndexMap[index])[kSCT];
              }
              else {
                 ++missing_truth_particle;
              }
           }
        }
      }
    } // Loop over SCT clusters
   
    for (const auto *const pix : *pixelClusters) {
      const xAOD::TrackMeasurementValidation* pixCluster = pix;
      static const SG::AuxElement::ConstAccessor< std::vector<unsigned int> > truthIndexAcc("truth_index");
      if (truthIndexAcc.isAvailable(*pixCluster)) {
        const std::vector<unsigned int> &truth_indices = truthIndexAcc(*pixCluster);
        for (auto index : truth_indices) {
           if (index != std::numeric_limits<unsigned int>::max()) {
              if (index < truthIndexMap.size() && truthIndexMap[index] != std::numeric_limits<unsigned int>::max()) {
                 ++tp_clustercount.at(truthIndexMap[index])[kPixel];
              }
              else {
                 ++missing_truth_particle;
              }
           }
        }
      }
    } // Loop over PIX clusters
  }
  m_nMissingTruthParticles += missing_truth_particle;
 
  if (not float_decor.empty()) {
     SG::ReadDecorHandle<xAOD::EventInfo, float> beamPosX(m_beamSpotDecoKey[0], ctx);
     SG::ReadDecorHandle<xAOD::EventInfo, float> beamPosY(m_beamSpotDecoKey[1], ctx);
     SG::ReadDecorHandle<xAOD::EventInfo, float> beamPosZ(m_beamSpotDecoKey[2], ctx);
     Amg::Vector3D beamPos = Amg::Vector3D(beamPosX(0), beamPosY(0), beamPosZ(0));
 
     if ( m_truthSelectionTool.get() ) {
        CutFlow tmp_cut_flow(m_truthSelectionTool->nCuts());
        for (const xAOD::TruthParticle *truth_particle : *ptruth) {
           auto passed = m_truthSelectionTool->accept(truth_particle);
           tmp_cut_flow.update( passed.missingCuts() );
           if (not passed) continue;
           decorateTruth(*truth_particle, float_decor, beamPos, tp_clustercount);
        }
        std::lock_guard<std::mutex> lock(m_mutex);
        m_cutFlow.merge(std::move(tmp_cut_flow));
     }
     else {
        for (const xAOD::TruthParticle *truth_particle : *ptruth) {
           decorateTruth(*truth_particle, float_decor, beamPos, tp_clustercount);
        }
     }
  }
 
  if (!decorateTruthTime(float_decor)) {
   return StatusCode::FAILURE;
  }
 
  return StatusCode::SUCCESS;
}
 
bool
InDetPhysValTruthDecoratorAlg::decorateTruth(const xAOD::TruthParticle& particle,
                                             std::vector<IDPVM::OptionalDecoration<xAOD::TruthParticleContainer, float> > &float_decor,
                                             const Amg::Vector3D& beamPos,
                                             const std::vector<std::array<uint16_t,kNClusterTypes> > &counts) const {
  ATH_MSG_VERBOSE("Decorate truth with d0 etc");
  if (particle.isNeutral()) {
    return false;
  }
  const EventContext& ctx = Gaudi::Hive::currentContext();
  const Amg::Vector3D momentum(particle.px(), particle.py(), particle.pz());
  const int pid(particle.pdgId());
  double charge = particle.charge();
 
  if (std::isnan(charge)) {
    ATH_MSG_DEBUG("charge not found on particle with pid " << pid);
    return false;
  }
   
  // @TODO float?
  float nSiHits = std::accumulate(counts.at(particle.index()).begin(), counts[particle.index()].end(), 0u);
 
  IDPVM::decorateOrRejectQuietly(particle,float_decor[kDecorNSilHits], nSiHits);
 
  const xAOD::TruthVertex* ptruthVertex(nullptr);
  try{
    ptruthVertex = particle.prodVtx();
  } catch (const std::exception& e) {
    if (not m_errorEmitted) {
      ATH_MSG_WARNING("A non existent production vertex was requested in calculating the track parameters d0 etc");
    }
    m_errorEmitted = true;
    return false;
  }
  if (!ptruthVertex) {
    ATH_MSG_DEBUG("A production vertex pointer was retrieved, but it is NULL");
    return false;
  }
  const auto xPos = ptruthVertex->x();
  const auto yPos = ptruthVertex->y();
  const auto z_truth = ptruthVertex->z();
  const Amg::Vector3D position(xPos, yPos, z_truth);
  const float prodR_truth = std::sqrt(xPos * xPos + yPos * yPos);
  // delete ptruthVertex;ptruthVertex=0;
  const Trk::CurvilinearParameters cParameters(position, momentum, charge);
 
  Trk::PerigeeSurface persf(beamPos);
 
  std::unique_ptr<const Trk::TrackParameters> tP ( m_extrapolator->extrapolate(ctx,
                                                                               cParameters, 
                                                                               persf, Trk::anyDirection, false) );
  if (tP) {
    float d0_truth = tP->parameters()[Trk::d0];
    float theta_truth = tP->parameters()[Trk::theta];
    float z0_truth = tP->parameters()[Trk::z0];
    float phi_truth = tP->parameters()[Trk::phi];
    float qOverP_truth = tP->parameters()[Trk::qOverP]; // P or Pt ??
    float z0st_truth = z0_truth * std::sin(theta_truth);
 
    // 'safeDecorator' used to prevent a crash in case of adding something which pre-exists.
    // behaviour chosen is to reject quietly
    IDPVM::decorateOrRejectQuietly(particle,float_decor[kDecorD0],d0_truth);
    IDPVM::decorateOrRejectQuietly(particle,float_decor[kDecorZ0],z0_truth);
    IDPVM::decorateOrRejectQuietly(particle,float_decor[kDecorPhi],phi_truth);
    IDPVM::decorateOrRejectQuietly(particle,float_decor[kDecorTheta],theta_truth);
    IDPVM::decorateOrRejectQuietly(particle,float_decor[kDecorZ0st],z0st_truth);
    IDPVM::decorateOrRejectQuietly(particle,float_decor[kDecorQOverP],qOverP_truth);
    IDPVM::decorateOrRejectQuietly(particle,float_decor[kDecorProdR],prodR_truth);
    IDPVM::decorateOrRejectQuietly(particle,float_decor[kDecorProdZ],z_truth);
 
    return true;
  } else {
    ATH_MSG_DEBUG("The TrackParameters pointer for this TruthParticle is NULL");
    return false;
  }
}
 
bool
InDetPhysValTruthDecoratorAlg::decorateTruthTime(std::vector<IDPVM::OptionalDecoration<xAOD::TruthParticleContainer, float>>& float_decor) const {
 
  const EventContext& ctx = Gaudi::Hive::currentContext();
 
  const xAOD::TruthVertex* truthVtx = nullptr;
  float truthTime;
 
  // First HS event
  if (!m_truthEventName.key().empty()) {
    ATH_MSG_VERBOSE("Getting TruthEventContainer");
    SG::ReadHandle<xAOD::TruthEventContainer> truthEventContainer(m_truthEventName, ctx);
    if (!truthEventContainer.isPresent()) {
      ATH_MSG_WARNING("TruthEventContainer name was specified, but no container is present");
      return true;
    }
    const xAOD::TruthEvent* event = (truthEventContainer.isValid()) ? truthEventContainer->at(0) : nullptr;
    if (event) {
      truthVtx = event->signalProcessVertex();
      truthTime = (truthVtx) ? truthVtx->t() / Gaudi::Units::c_light : -9999.;
      for (const auto& link : event->truthParticleLinks()) {
        if (link.isValid()) {
          IDPVM::decorateOrRejectQuietly(**link, float_decor[kDecorTime], truthTime);
        }
      }
    }
    else {
      ATH_MSG_ERROR("No valid TruthEvent!");
      return false;
    }
  }
 
   // Then PU events
  if (!m_truthPileupEventName.key().empty()) {
    ATH_MSG_VERBOSE("Getting TruthPileupEventContainer");
    SG::ReadHandle<xAOD::TruthPileupEventContainer> truthPileupEventContainer(m_truthPileupEventName, ctx);
    if (!truthPileupEventContainer.isPresent()) {
      ATH_MSG_WARNING("TruthPileupEventContainer name was specified, but no container is present");
      return true;
    }
    if (truthPileupEventContainer.isValid()) {
      for (const auto event : *truthPileupEventContainer) {
        truthVtx = nullptr;
        for (std::size_t i = 0; i < event->nTruthVertices(); i++) {
          truthVtx = event->truthVertex(i);
          if (truthVtx) {
            break;
          }
        }
        truthTime = (truthVtx) ? truthVtx->t() / Gaudi::Units::c_light: -9999.;
        for (const auto& link : event->truthParticleLinks()) {
          if (link.isValid()) {
            IDPVM::decorateOrRejectQuietly(**link, float_decor[kDecorTime], truthTime);
          }
        }
      }
    }
    else {
      ATH_MSG_ERROR("TruthPileupEventContainer is invalid!");
      return false;
    }
  }
 
  return true;
 
}