DecoratePromptLeptonImproved.cxx

Go to the documentation of this file.

/*
  Copyright (C) 2002-2025 CERN for the benefit of the ATLAS collaboration
*/
 
// Local
#include "LeptonTaggers/DecoratePromptLeptonImproved.h"
#include "LeptonTaggers/PromptUtils.h"
 
// ROOT
#include "TMVA/Config.h"
#include "TH1.h"
 
// C/C++
#include <cmath>
#include <iostream>
#include <sstream>
#include <sys/stat.h>
 
 
namespace {
 
 
// Accessor constants.
static const SG::ConstAccessor<float> accessCalIsolation30 ("topoetcone30");
static const SG::ConstAccessor<float> accessTrackIsolation30 ("ptvarcone30");
static const SG::ConstAccessor<float> accessTrackIsolation30TTVA ("ptvarcone30_Nonprompt_All_MaxWeightTTVA_pt500");
static const SG::ConstAccessor<float> accessMuonCalE ("calE");
static const SG::ConstAccessor<float> accessMuonParamEnergyLoss ("ParamEnergyLoss");
 
 
} // anonymous namespace
 
 
//======================================================================================================
Prompt::DecoratePromptLeptonImproved::DecoratePromptLeptonImproved(const std::string& name, ISvcLocator* pSvcLocator):
  AthAlgorithm(name, pSvcLocator),
  m_BDTVarKey(Def::NONE)
{}
 
//=============================================================================
StatusCode Prompt::DecoratePromptLeptonImproved::initialize()
{
  ATH_MSG_DEBUG("Initializing DecoratePromptLeptonImproved...");
  ATH_MSG_DEBUG("m_leptonsName = " << m_leptonsName);
  if (m_printTime)
  {
    //
    // Reset timers
    //
    m_timerAll .Reset();
    m_timerExec.Reset();
    m_timerMuon.Reset();
    m_timerElec.Reset();
 
    //
    // Start full timer
    //
    m_timerAll.Start();
  }
 
  ATH_CHECK(m_trackJetsKey.initialize());
  ATH_CHECK(m_primaryVertexKey.initialize());
  ATH_CHECK(m_clusterContainerKey.initialize());
 
  ATH_MSG_DEBUG("Initializing " << m_electronsKey);
 
  ATH_CHECK(m_electronsKey.initialize(m_leptonsName == "Electrons"));
  ATH_CHECK(m_muonsKey.initialize(m_leptonsName == "Muons"));
 
  ATH_MSG_DEBUG("Number of int vars to read: " << m_stringIntVars.size());
  ATH_MSG_DEBUG("Number of float vars to read: " << m_stringFloatVars.size());
 
  // Setup variable holder
  m_vars = std::make_unique<Prompt::VarHolder>();
 
  //
  // Read vector<string> vars into vector<Var> and append
  //
  m_intVars   = m_vars->readVectorVars(m_stringIntVars);
  m_floatVars = m_vars->readVectorVars(m_stringFloatVars);
 
  m_allVars.insert(m_allVars.end(), m_intVars  .begin(), m_intVars  .end());
  m_allVars.insert(m_allVars.end(), m_floatVars.begin(), m_floatVars.end());
 
  m_varTMVA.resize(m_allVars.size());
 
 
  //
  // Get key for recording BDT output
  //
  m_BDTVarKey = m_vars->registerDynamicVar(m_BDTName);
  
  if(m_BDTVarKey == Def::NONE) {
    ATH_MSG_ERROR("Failed to create key for BDT name=" << m_BDTName);
    return StatusCode::FAILURE;
  }
 
  //
  // Fill decorator maps
  //
  ATH_CHECK(initializeDecorators());
 
  //
  // Initialize const accessors
  //
  initializeConstAccessors();
 
  ATH_MSG_DEBUG("Initialized DecoratePromptLeptonImproved.");
 
  //
  // Initialize TMVA Reader
  //
  ATH_CHECK(initializeTMVAReader());
  
  ATH_MSG_DEBUG("Initialized TMVA Reader.");
 
  return StatusCode::SUCCESS;
}
 
//=============================================================================
StatusCode Prompt::DecoratePromptLeptonImproved::finalize()
{
  if(m_printTime) {
    //
    // Print full time stopwatch
    //
    m_timerAll.Stop();
 
    ATH_MSG_INFO("Real time: " << m_timerAll.RealTime() << "\t CPU time: " << m_timerAll.CpuTime());
 
    ATH_MSG_INFO("Execute time: " << PrintResetStopWatch(m_timerExec));
    ATH_MSG_INFO("Muon    time: " << PrintResetStopWatch(m_timerMuon));
    ATH_MSG_INFO("Elec    time: " << PrintResetStopWatch(m_timerElec));
  }
 
  return StatusCode::SUCCESS;
}
 
 
//=============================================================================
StatusCode Prompt::DecoratePromptLeptonImproved::execute(const EventContext& ctx)
{
  //
  // Start execute timer
  //
  TimerScopeHelper timer(m_timerExec);
 
 
  //
  // Retrieve containers from evtStore
  //
  SG::ReadHandle<xAOD::JetContainer> trackJets(m_trackJetsKey, ctx);
  SG::ReadHandle<xAOD::VertexContainer> vertices(m_primaryVertexKey, ctx);
  SG::ReadHandle<xAOD::CaloClusterContainer> clusters(m_clusterContainerKey, ctx);
 
  ATH_MSG_DEBUG("======================================="
    << "\n\t\t\t  Size of vertex    container: " << vertices ->size()
    << "\n\t\t\t  Size of track jet container: " << trackJets->size()
    << "\n-----------------------------------------------------------------");
 
  //
  // Find default Primary Vertex
  //
  const xAOD::Vertex *primaryVertex = nullptr;
 
  for(const xAOD::Vertex *vertex: *vertices) {
    if(vertex->vertexType() == xAOD::VxType::PriVtx) {
      primaryVertex = vertex;
      break;
    }
  }
 
  std::vector<decoratorFloatH_t> floatDecors;
  floatDecors.reserve (m_floatDecorHandleKeys.size());
  for (const SG::WriteDecorHandleKey<xAOD::IParticleContainer>& k : m_floatDecorHandleKeys) {
    floatDecors.emplace_back (k, ctx);
  }
 
  std::vector<decoratorShortH_t> shortDecors;
  shortDecors.reserve (m_shortDecorHandleKeys.size());
  for (const SG::WriteDecorHandleKey<xAOD::IParticleContainer>& k : m_shortDecorHandleKeys) {
    shortDecors.emplace_back (k, ctx);
  }
 
  if(m_leptonsName == "Electrons") {
    //
    // Process electrons
    //
    ATH_MSG_DEBUG("Reading " << m_electronsKey);
    SG::ReadHandle<xAOD::ElectronContainer> electrons(m_electronsKey, ctx);
 
    for(const xAOD::Electron *elec: *electrons) {
      decorateElec(*elec, *trackJets, *clusters, primaryVertex, floatDecors, shortDecors);
    }
  } else if(m_leptonsName == "Muons") {
    //
    // Process muons
    //
    ATH_MSG_DEBUG("Reading " << m_muonsKey);
    SG::ReadHandle<xAOD::MuonContainer> muons(m_muonsKey, ctx);
 
    for(const xAOD::Muon *muon: *muons) {
      decorateMuon(*muon, *trackJets, primaryVertex, floatDecors, shortDecors);
      ATH_MSG_DEBUG("Muon decorated");
    }
  } else {
    ATH_MSG_ERROR("Must specify Electrons or Muons");
    return StatusCode::FAILURE;
  }
 
  return StatusCode::SUCCESS;
}
 
 
//=============================================================================
bool Prompt::DecoratePromptLeptonImproved::initializeTMVAReader()
{   
  //
  // Make new instance of TMVA Reader and add variables
  //
  m_TMVAReader = std::make_unique<TMVA::Reader>();
 
  for(unsigned i = 0; i < m_allVars.size(); ++i) { 
    m_TMVAReader->AddVariable(m_vars->asStr(m_allVars.at(i)), &m_varTMVA[i]);
  }
  
  //
  // Get path to xml training file
  //
  std::string fullPathToFile = PathResolverFindCalibFile("JetTagNonPromptLepton/" 
                             + m_configFileVersion 
                             + "/TMVAClassification_" + m_methodTitleMVA + ".weights.xml");
 
  if(!m_configPathOverride.empty()) {
    ATH_MSG_INFO("Override path resolver result");
    fullPathToFile = m_configPathOverride;
  }
 
  ATH_MSG_INFO("TMVA configuration file: " + fullPathToFile);
  ATH_MSG_INFO("TMVA method name: " + m_methodTitleMVA);
 
  //
  // Book an instance of BDT reader if path is available
  //
  if(fullPathToFile == "") {
    ATH_MSG_ERROR("Could not find path to xml training file");
    return false;
  }
  else {
    m_TMVAReader->BookMVA(m_methodTitleMVA.toString(), fullPathToFile);
  }
 
  return true;
}
 
 
//=============================================================================
StatusCode Prompt::DecoratePromptLeptonImproved::initializeDecorators()
{
  std::string inputContName;
  if (m_leptonsName == "Electrons") {
    inputContName = m_electronsKey.key();
  }
  else if (m_leptonsName == "Muons") {
    inputContName = m_muonsKey.key();
  }
  else {
    ATH_MSG_ERROR ("LeptonContainerName was " << m_leptonsName <<
                   "; must be either Electrons or Muons");
    return StatusCode::FAILURE;
  }
 
  DecorMap_t vetoedVars;
  for(const std::string &vvar: m_vetoDecoratorFloatVars) {
    const Def::Var vkey = m_vars->registerDynamicVar(vvar);
 
    if(vkey == Def::NONE) {
      ATH_MSG_ERROR("Failed to create key for variable name=" << vvar);
      return StatusCode::FAILURE;
    }
 
    vetoedVars.try_emplace (vkey, 0);
  }
  for(const std::string &vvar: m_vetoDecoratorShortVars) {
    const Def::Var vkey = m_vars->registerDynamicVar(vvar);
 
    if(vkey == Def::NONE) {
      ATH_MSG_ERROR("Failed to create key for variable name=" << vvar);
      return StatusCode::FAILURE;
    }
 
    vetoedVars.try_emplace (vkey, 0);
  }
 
  //
  // Fill short variable map
  //
  for(Prompt::Def::Var &var: m_intVars) {
    if (!vetoedVars.contains (var)) {
      m_shortDecorMap.try_emplace (var, m_shortDecorHandleKeys.size());
      m_shortDecorHandleKeys.push_back (SG::makeContDecorKey (inputContName, m_inputVarDecoratePrefix + m_vars->asStr(var)));
    }
  }
 
  //
  // Fill float variable map
  //
  m_floatDecorHandleKeys.clear();
  for(Prompt::Def::Var &var: m_floatVars) {
    if (!vetoedVars.contains (var)) {
      m_floatDecorMap.try_emplace (var, m_floatDecorHandleKeys.size());
      m_floatDecorHandleKeys.push_back (SG::makeContDecorKey (inputContName, m_inputVarDecoratePrefix + m_vars->asStr(var)));
    }
  }
 
  //
  // Fill additional variables
  //
  if (!vetoedVars.contains (m_BDTVarKey)) {
    m_floatDecorMap.try_emplace (m_BDTVarKey, m_floatDecorHandleKeys.size());
    m_floatDecorHandleKeys.push_back (SG::makeContDecorKey (inputContName, m_BDTName));
  }
 
  for(const std::string &evar: m_extraDecoratorFloatVars) {
    const Def::Var ekey = m_vars->registerDynamicVar(evar);
 
    if(ekey == Def::NONE) {
      ATH_MSG_ERROR("Failed to create key for variable name=" << evar);
      return StatusCode::FAILURE;
    }
 
    if (!vetoedVars.contains (ekey)) {
      m_floatDecorMap.try_emplace (ekey, m_floatDecorHandleKeys.size());
      m_floatDecorHandleKeys.push_back (SG::makeContDecorKey (inputContName, m_inputVarDecoratePrefix + evar));
    }
  }
 
  for(const std::string &evar: m_extraDecoratorShortVars) {
    const Def::Var ekey = m_vars->registerDynamicVar(evar);
 
    if(ekey == Def::NONE) {
      ATH_MSG_ERROR("Failed to create key for variable name=" << evar);
      return StatusCode::FAILURE;
    }
 
    if (!vetoedVars.contains (ekey)) {
      m_shortDecorMap.try_emplace (ekey, m_shortDecorHandleKeys.size());
      m_shortDecorHandleKeys.push_back (SG::makeContDecorKey (inputContName, m_inputVarDecoratePrefix + evar));
    }
  }
 
  ATH_CHECK(m_shortDecorHandleKeys.initialize());
  ATH_CHECK(m_floatDecorHandleKeys.initialize());
 
  // Print decorator counts
  ATH_MSG_DEBUG("Added " << m_shortDecorHandleKeys.size() << " short decorators");
  ATH_MSG_DEBUG("Added " << m_floatDecorHandleKeys.size() << " float decorators");
 
  // Instantiate MVA X bin
  if(m_leptonPtBinsVector.size() < 2) {
    ATH_MSG_ERROR("Invalid PtBins size=" << m_leptonPtBinsVector.size());
    return StatusCode::FAILURE;
  }
 
  std::unique_ptr<double []> PtBins = std::make_unique<double []>(m_leptonPtBinsVector.size());
 
  for(unsigned i = 0; i < m_leptonPtBinsVector.size(); i++) {
    PtBins[i] = m_leptonPtBinsVector[i];
  }
 
  m_leptonPtBinHist = std::make_unique<TH1D>("PtBin", "PtBin", m_leptonPtBinsVector.size() - 1, PtBins.get());
 
  return StatusCode::SUCCESS;
}
 
//=============================================================================
void Prompt::DecoratePromptLeptonImproved::initializeConstAccessors()
{
  //
  // Instantiate isolation accessors
  //
  m_accessDeepSecondaryVertex.emplace (m_vertexLinkName);
 
  //
  // Instantiate accessors for RNN variables
  //
  for(const std::string &name: m_accessorRNNVars) {
    const Def::Var akey = m_vars->registerDynamicVar(name);
 
    if(m_accessRNNMap.insert(floatAccessorMap::value_type(akey, AccessFloat(name))).second) {
      ATH_MSG_DEBUG("Add float RNN accessor: " << name);
    }
    else {
      ATH_MSG_WARNING("Skip duplicate float accessor: " << name);
    }
  }
}
 
//=============================================================================
void Prompt::DecoratePromptLeptonImproved::decorateElec(
  const xAOD::Electron &electron,
  const xAOD::JetContainer &trackJets,
  const xAOD::CaloClusterContainer &clusters,
  const xAOD::Vertex *primaryVertex,
  std::vector<decoratorFloatH_t>& floatDecors,
  std::vector<decoratorShortH_t>& shortDecors
)
{
  //
  // Find nearest track jet to electron
  //
  TimerScopeHelper timer(m_timerElec);
 
  Prompt::VarHolder vars;
 
  std::pair<double, const xAOD::Jet*> match = findTrackJet(electron, trackJets);
 
  if(match.second) {
    //
    // Get muon calorimeter energy variable, RNN and secondary vertex variables
    //
    getElectronAnpVariables(electron, clusters, vars, primaryVertex);
 
    //
    // Get mutual variables, passing track as argument
    //
    getMutualVariables(electron, *match.second, electron.trackParticle(), vars);
 
    //
    // Pass variables to TMVA
    //
    addVarsToTMVA(vars);
  }
  else {
    //
    // Decorate electron with default values
    //
    fillVarDefault(vars);
 
    ATH_MSG_DEBUG("No track jet found near to electron");
  }
 
  //
  // Decorate electron with input vars and BDT weight
  //
  decorateAuxLepton(electron, vars, floatDecors, shortDecors);
}
 
 
//=============================================================================
void Prompt::DecoratePromptLeptonImproved::decorateMuon(
  const xAOD::Muon         &muon,
  const xAOD::JetContainer &trackJets,
  const xAOD::Vertex *primaryVertex,
  std::vector<decoratorFloatH_t>& floatDecors,
  std::vector<decoratorShortH_t>& shortDecors
)
{
  //
  // Find nearest track jet to muon
  //
  TimerScopeHelper timer(m_timerMuon);
 
  Prompt::VarHolder vars;
 
  std::pair<double, const xAOD::Jet*> match = findTrackJet(muon, trackJets);
 
  if(match.second) {
    //
    // Get muon calorimeter energy variable, RNN and secondary vertex variables
    //
    getMuonAnpVariables(muon, vars, primaryVertex);
 
    //
    // Get mutual variables, passing track as argument
    //
    getMutualVariables(muon, *match.second, muon.trackParticle(xAOD::Muon::TrackParticleType::Primary), vars);
 
    //
    // Add variables to TMVA Reader
    //
    addVarsToTMVA(vars);
  }
  else {
    //
    // Decorate muon with default values
    //
    fillVarDefault(vars);
 
    ATH_MSG_DEBUG("No track jet found near to muon");
  }
 
  //
  // Decorate muon with input vars and BDT weight
  //
  decorateAuxLepton(muon, vars, floatDecors, shortDecors);
}
 
//=============================================================================
void Prompt::DecoratePromptLeptonImproved::getElectronAnpVariables(
  const xAOD::Electron             &elec,
  const xAOD::CaloClusterContainer &clusters,
  Prompt::VarHolder                &vars,
  const xAOD::Vertex               *primaryVertex
)
{
  //
  // Get Muon variables - calorimeter
  //
  double sumCoreEt_large = 0.0, CaloClusterSumEtRel = 0.0;
 
  if(elec.caloCluster()) {
    const double elec_calEta = elec.caloCluster()->eta();
    const double elec_calPhi = elec.caloCluster()->phi();
 
    for(const xAOD::CaloCluster *cluster: clusters) {
      const double deta =                      elec_calEta - cluster->eta();
      const double dphi = TVector2::Phi_mpi_pi(elec_calPhi - cluster->phi());
      const double dr   = std::sqrt(deta*deta + dphi*dphi);
 
      if(dr < m_elecMinCalErelConeSize) {
        sumCoreEt_large += cluster->pt();
      }
    }
  }
 
  if(elec.pt() > 0.0) CaloClusterSumEtRel = sumCoreEt_large/elec.pt();
 
  vars.addVar(Prompt::Def::CaloClusterSumEtRel, CaloClusterSumEtRel);
 
  //
  // Get lepton isolation variables
  //
  const double Topoetcone30rel = accessIsolation(accessCalIsolation30, elec);
  const double Ptvarcone30rel  = accessIsolation(accessTrackIsolation30, elec);
 
  vars.addVar(Prompt::Def::Topoetcone30rel, Topoetcone30rel);
  vars.addVar(Prompt::Def::Ptvarcone30rel,  Ptvarcone30rel);
 
  //
  // Get secondary vertex variable
  //
  std::vector<double> goodVertexNdistLong;
 
  if(m_accessDeepSecondaryVertex->isAvailable(elec)) {
    std::vector<ElementLink<xAOD::VertexContainer> > vtxLinks = (*m_accessDeepSecondaryVertex)(elec);
 
    for(ElementLink<xAOD::VertexContainer> &vtxLink: vtxLinks) {
      if(!vtxLink.isValid()) {
        ATH_MSG_WARNING("VertexContainer : invalid link");
        continue;
      }
 
      const xAOD::Vertex *vtx = *vtxLink;
 
      const double fitProb = Prompt::getVertexFitProb(vtx);
 
      if(fitProb < m_vertexMinChiSquaredProb) {
        continue;
      }
 
      const double theta = std::acos(getVertexCosThetaWithLepDir(elec, vtx, primaryVertex));
 
      if      (theta < m_vertexMinThetaBarrElec && std::fabs(elec.eta()) <= m_vertexBarrEcapAbsEtaAt) continue;
      else if (theta < m_vertexMinThetaEcapElec && std::fabs(elec.eta()) >  m_vertexBarrEcapAbsEtaAt) continue;
 
      const double vertex_ndist_long = getVertexLongitudinalNormDist(elec, vtx, primaryVertex);
 
      goodVertexNdistLong.push_back(vertex_ndist_long);
    }
  }
  else {
    ATH_MSG_WARNING("VertexContainer : " << m_vertexLinkName << " not found for the electron");
  }
 
  double best_vertex_ndist_long = 0.0;
 
  if(goodVertexNdistLong.size() > 0) {
    std::sort(goodVertexNdistLong.begin(), goodVertexNdistLong.end());
    best_vertex_ndist_long = goodVertexNdistLong.back();
  }
 
  vars.addVar(Prompt::Def::CandVertex_normDistToPriVtxLongitudinalBest_ThetaCutVtx, best_vertex_ndist_long);
  vars.addVar(Prompt::Def::CandVertex_NPassVtx,                         goodVertexNdistLong.size());
}
 
//=============================================================================
void Prompt::DecoratePromptLeptonImproved::getMuonAnpVariables(
  const xAOD::Muon   &muon,
  Prompt::VarHolder                &vars,
  const xAOD::Vertex *primaryVertex
) const
{
  //
  // Get Muon variables - calorimeter
  //
  double calE = -99.0, peloss = -99.0, caloClusterERel = -99.0;
 
  if(const xAOD::CaloCluster* cluster = muon.cluster()) {
    if(accessMuonCalE.isAvailable(*cluster) && accessMuonParamEnergyLoss.isAvailable(muon)) {
      calE   = accessMuonCalE(*cluster);
      peloss = accessMuonParamEnergyLoss(muon);
 
      caloClusterERel = calE/peloss;
    }
    else {
      ATH_MSG_WARNING("Muon calE or ParamEnergyLoss not found in auxiliary store");
    }
  }
 
  vars.addVar(Prompt::Def::CaloClusterERel, caloClusterERel);
 
  //
  // Get lepton isolation variables
  //
  const double Topoetcone30rel = accessIsolation(accessCalIsolation30, muon);
  const double ptvarcone30TightTTVAPt500rel = accessIsolation(accessTrackIsolation30TTVA, muon);
 
  vars.addVar(Prompt::Def::Topoetcone30rel,              Topoetcone30rel);
  vars.addVar(Prompt::Def::Ptvarcone30_TightTTVA_pt500rel, ptvarcone30TightTTVAPt500rel);
 
  //
  // Get Muon Secondary Vertex variable
  //
  std::vector<double> goodVertexNdistLong;
 
  if(m_accessDeepSecondaryVertex->isAvailable(muon)) {
    std::vector<ElementLink<xAOD::VertexContainer> > vtxLinks = (*m_accessDeepSecondaryVertex)(muon);
    goodVertexNdistLong.reserve(vtxLinks.size());
 
    for(ElementLink<xAOD::VertexContainer> &vtxLink: vtxLinks) {
      if(!vtxLink.isValid()) {
        ATH_MSG_WARNING("VertexContainer : invalid link");
        continue;
      }
 
      const xAOD::Vertex *vtx = *vtxLink;
 
      const double fitProb = Prompt::getVertexFitProb(vtx);
 
      if(fitProb > m_vertexMinChiSquaredProb) {
        const double vertex_ndist_long = getVertexLongitudinalNormDist(muon, vtx, primaryVertex);
 
        goodVertexNdistLong.push_back(vertex_ndist_long);
      }
    }
  }
  else {
    ATH_MSG_WARNING("VertexContainer : " << m_vertexLinkName << " not found for the muon");
  }
 
  double best_vertex_ndist_long = 0.0;
 
  if(goodVertexNdistLong.size() > 0) {
    std::sort(goodVertexNdistLong.begin(), goodVertexNdistLong.end());
    best_vertex_ndist_long = goodVertexNdistLong.back();
  }
 
  vars.addVar(Prompt::Def::CandVertex_normDistToPriVtxLongitudinalBest, best_vertex_ndist_long);
  vars.addVar(Prompt::Def::CandVertex_NPassVtx,                         goodVertexNdistLong.size());
}
 
//=============================================================================
void Prompt::DecoratePromptLeptonImproved::getMutualVariables(
  const xAOD::IParticle     &particle,
  const xAOD::Jet           &track_jet,
  const xAOD::TrackParticle *track,
  Prompt::VarHolder         &vars
) const
{
  //
  // Add lepton - jet variables to VarHolder
  //
  double PtFrac = -99.;
  double PtRel  = -99.;
 
  if(particle.pt() > 0.0 && track_jet.pt() > 0.0) {
 
    if(track) {
      PtFrac = track->pt()/track_jet.pt();
    }
 
    const double angle = particle.p4().Vect().Angle(track_jet.p4().Vect());
 
    PtRel = particle.pt() * std::sin(angle);
  }
 
  //
  // Add vars to VarHolder
  //
  vars.addVar(Prompt::Def::PtFrac,         PtFrac);
  vars.addVar(Prompt::Def::PtRel,          PtRel);
  vars.addVar(Prompt::Def::DRlj,           track_jet.p4().DeltaR(particle.p4()));
  vars.addVar(Prompt::Def::TrackJetNTrack, track_jet.getConstituents().size());
 
  //
  // Get RNN variables
  //
  for(const floatAccessorMap::value_type &acc: m_accessRNNMap) {
    if(acc.second.isAvailable(particle)) {
      vars.addVar(acc.first, acc.second(particle));
    }
    else {
      ATH_MSG_WARNING("LeptonTagger RNN not found in auxiliary store for variable=" << vars.asStr(acc.first));
    }
  }
 
  //
  // Get lepton variables - pT X bin
  //
  const double lepPt = particle.pt();
 
  const double xmax = m_leptonPtBinHist->GetXaxis()->GetXmax();
  const double xmin = m_leptonPtBinHist->GetXaxis()->GetXmin();
 
  int curr_bin = 0;
 
  if(xmin < lepPt && lepPt < xmax) {
    curr_bin = m_leptonPtBinHist->FindBin(lepPt);
  }
  else if (!(lepPt < xmax)) {
    curr_bin = m_leptonPtBinHist->GetNbinsX();
  }
  else if (!(lepPt > xmin)) {
    curr_bin = 1;
  }
 
  vars.addVar(Prompt::Def::MVAXBin, curr_bin);
  vars.addVar(Prompt::Def::RawPt,   lepPt);
 
  ATH_MSG_DEBUG("getMutualVariables - lepPt = " << lepPt << ", MVAXBin = " << curr_bin);
}
 
//=============================================================================
float Prompt::DecoratePromptLeptonImproved::accessIsolation(
  const AccessFloat           &isoAccessor,
  const xAOD::IParticle &particle
) const {
  double isolation = -99., isolationrel = -99.;
 
  if(isoAccessor.isAvailable(particle)) {
    isolation = isoAccessor(particle);
  }
  else {
    ATH_MSG_WARNING("Lepton isolation not found in auxiliary store");
  }
 
  if(particle.pt() > 0.0) {
    isolationrel = isolation / particle.pt();
  }
 
  return isolationrel;
}
 
//=============================================================================
void Prompt::DecoratePromptLeptonImproved::addVarsToTMVA(Prompt::VarHolder &vars)
{
  //
  // Add variables to TMVA reader
  //
  for(unsigned i = 0; i < m_allVars.size(); ++i) {  
    m_varTMVA[i] = 0.0;
 
    if(!vars.getVar(m_allVars.at(i), m_varTMVA[i])) {
      ATH_MSG_WARNING("Missing input variable: " << m_vars->asStr(m_allVars.at(i)));   
    }
  }
 
  //
  // Decorate lepton with classifier response, if goodJet
  //
  float bdt_weight = m_TMVAReader->EvaluateMVA(m_methodTitleMVA.toString());
 
  if(m_BDTVarKey != Def::NONE) {
    vars.addVar(m_BDTVarKey, bdt_weight);
  }
  else {
    ATH_MSG_WARNING("addVarsToTMVA - invalid Def::Var key for " << m_BDTName);
  }
}
 
//=============================================================================
void Prompt::DecoratePromptLeptonImproved::fillVarDefault(Prompt::VarHolder &vars) const
{
  //
  // Add default values to VarHolder
  //
  for(const DecorMap_t::value_type &dec: m_floatDecorMap) {
    vars.addVar(dec.first, -99.0);
  }
 
  for(const DecorMap_t::value_type &dec: m_shortDecorMap) {
    vars.addVar(dec.first, -99.0);
  }
}
 
//=============================================================================
void Prompt::DecoratePromptLeptonImproved::decorateAuxLepton(
  const xAOD::IParticle &particle,
  Prompt::VarHolder &vars,
  std::vector<decoratorFloatH_t>& floatDecors,
  std::vector<decoratorShortH_t>& shortDecors
) const
{
  //
  // Decorate lepton with input short variables
  //
  for(const DecorMap_t::value_type &dec: m_shortDecorMap) {
    double val = 0.0;
 
    if(vars.getVar(dec.first, val)) {
      shortDecors.at(dec.second)(particle) = static_cast<short>(val);
 
      ATH_MSG_DEBUG("Short variable: " << vars.asStr(dec.first) << " = " << val);
    }
    else {
      ATH_MSG_WARNING("Short variable " << vars.asStr(dec.first) << " not decorated to lepton");
    }
  }
 
  //
  // Decorate lepton with input float variables
  //
  for(const DecorMap_t::value_type &dec: m_floatDecorMap) {
    double val = 0.0;
 
    if(vars.getVar(dec.first, val)) {
      floatDecors.at(dec.second)(particle) = val;
 
      ATH_MSG_DEBUG("Float variable: " << vars.asStr(dec.first) << " = " << val);
    }
    else {
      ATH_MSG_WARNING("Float variable " << vars.asStr(dec.first) << " not decorated to lepton");
    }
  }
}
 
//=============================================================================
template<class T> std::pair<double, const xAOD::Jet*> Prompt::DecoratePromptLeptonImproved::findTrackJet(const T &part,
                          const xAOD::JetContainer &jets)
{
  //
  // Find nearest track jet and a return a pair of dR and xAOD::Jet*
  //
  const xAOD::Jet *minjet = 0;
  double           mindr  = 10.0;
  std::pair<double, const xAOD::Jet*> match(mindr, minjet);
 
  for(const xAOD::Jet* jet: jets) {
    const double dr = part.p4().DeltaR(jet->p4());
 
    if(!minjet || dr < mindr) {
      mindr  = dr;
      minjet = jet;
      match = std::make_pair(mindr, minjet);
    }
  }
 
 
  if(minjet && mindr < m_maxLepTrackJetDR) {
    return match;
  }
  
  minjet = 0;
  return std::make_pair(10., minjet);
}
 
//=============================================================================
double Prompt::DecoratePromptLeptonImproved::getVertexLongitudinalNormDist(
  const xAOD::IParticle &lepton,
  const xAOD::Vertex    *secondaryVertex,
  const xAOD::Vertex    *primaryVertex
) const
{
  //
  // get the Longitudinal normalized distance between the secondary vertex and primary vertex
  //
  if(!secondaryVertex || !primaryVertex) {
    ATH_MSG_WARNING("getVertexLongitudinalNormDist - invalid pointer of lepton/secondaryVertex/primaryVertex");
    return 0.0;
  }
 
 
  float normDist_SVPV  = 0.0;
 
  if(!Prompt::GetAuxVar(*secondaryVertex, normDist_SVPV,  "normDistToPriVtx")) {
    ATH_MSG_WARNING("getVertexLongitudinalNormDist - missing \"normDistToPriVtx\"");
  }
 
  double cos_theta = getVertexCosThetaWithLepDir(lepton, secondaryVertex, primaryVertex);
 
  return normDist_SVPV*cos_theta;
}
 
//=============================================================================
double Prompt::DecoratePromptLeptonImproved::getVertexCosThetaWithLepDir(const xAOD::IParticle &lepton,
                                                                         const xAOD::Vertex    *secondaryVertex,
                                                                         const xAOD::Vertex    *primaryVertex) const
{
  //
  // get the Longitudinal nomalized distance between the secondary vertex and primary vertex
  //
  if(!secondaryVertex || !primaryVertex) {
    ATH_MSG_WARNING("GetVertexThetaWithLepDir - invalid pointer of lepton/secondaryVertex/primaryVertex");
    return 0.0;
  }
 
  const Amg::Vector3D sv_to_pv_v3 = secondaryVertex->position() - primaryVertex->position();
 
  const TVector3 sv_to_pv_t3 = TVector3(sv_to_pv_v3.x(), sv_to_pv_v3.y(), sv_to_pv_v3.z());
  const TVector3 lepton_dirt = lepton.p4().Vect();
 
  const double cos_theta = sv_to_pv_t3.Unit()*lepton_dirt.Unit();
 
  return cos_theta;
}