JSSWTopTaggerDNN.cxx

Go to the documentation of this file.

/*
  Copyright (C) 2002-2023 CERN for the benefit of the ATLAS collaboration
*/
 
#include "BoostedJetTaggers/JSSWTopTaggerDNN.h"
 
#include <fstream>
 
JSSWTopTaggerDNN::JSSWTopTaggerDNN( const std::string& name ) :
  JSSTaggerBase( name ),
  m_lwnn(nullptr)
{
 
}
 
StatusCode JSSWTopTaggerDNN::initialize() {
 
  ATH_MSG_INFO( "Initializing JSSWTopTaggerDNN tool" );
 
  m_ptGeV = true;
 
  m_useMassCut = true;
 
  m_useScoreCut = true;
 
  if ( ! m_configFile.empty() ) {
 
    ATH_CHECK( getConfigReader() );
 
    m_tagType = m_configReader.GetValue("TaggerType" ,"");
 
    m_kerasCalibArea = m_configReader.GetValue("CalibAreaKeras" ,"");
 
    m_kerasConfigFileName = m_configReader.GetValue("KerasConfigFile" ,"");
 
    m_kerasConfigOutputName = m_configReader.GetValue("KerasOutput" ,"");
 
    m_strMassCutLow  = m_configReader.GetValue("MassCutLow" ,"");
    m_strMassCutHigh = m_configReader.GetValue("MassCutHigh" ,"");
    m_strScoreCut    = m_configReader.GetValue("ScoreCut" ,"");
 
    m_jetPtMin = m_configReader.GetValue("pTCutLow", 350.0);
    m_jetPtMax = m_configReader.GetValue("pTCutHigh", 4000.0);
 
    m_decorationName = m_configReader.GetValue("DecorationName" ,"");
 
    m_calcSF = m_configReader.GetValue("CalcSF", false);
    if ( m_calcSF ) {
      m_weightDecorationName = m_configReader.GetValue("WeightDecorationName", "");
      m_weightFileName = m_configReader.GetValue("WeightFile", "");
      m_weightHistogramName = m_configReader.GetValue("WeightHistogramName", "");
      m_efficiencyHistogramName = m_configReader.GetValue("EfficiencyHistogramName", "");
      m_weightFlavors = m_configReader.GetValue("WeightFlavors", "");
 
      m_truthLabelName = m_configReader.GetValue("TruthLabelName" , "R10TruthLabel_R21Consolidated");
    }
 
    ATH_MSG_INFO( "Configurations Loaded :");
    ATH_MSG_INFO( "tagType               : " << m_tagType );
    ATH_MSG_INFO( "calibarea_keras       : " << m_kerasCalibArea );
    ATH_MSG_INFO( "kerasConfigFileName   : " << m_kerasConfigFileName );
    ATH_MSG_INFO( "kerasConfigOutputName : " << m_kerasConfigOutputName );
    ATH_MSG_INFO( "strMassCutLow         : " << m_strMassCutLow );
    ATH_MSG_INFO( "strMassCutHigh        : " << m_strMassCutHigh );
    ATH_MSG_INFO( "pTCutLow              : " << m_jetPtMin );
    ATH_MSG_INFO( "pTCutHigh             : " << m_jetPtMax );
    ATH_MSG_INFO( "strScoreCut           : " << m_strScoreCut );
    ATH_MSG_INFO( "decorationName        : " << m_decorationName );
    if ( m_calcSF ) {
      ATH_MSG_INFO( "weightDecorationName  : " << m_weightDecorationName );
      ATH_MSG_INFO( "weightFile            : " << m_weightFileName );
      ATH_MSG_INFO( "weightHistogramName   : " << m_weightHistogramName );
      ATH_MSG_INFO( "efficiencyHistogramName : "<<m_efficiencyHistogramName );
      ATH_MSG_INFO( "weightFlavors         : " << m_weightFlavors );
      ATH_MSG_INFO( "TruthLabelName        : " << m_truthLabelName );
    }
  }
  else { 
    if ( (m_kerasConfigFileName.empty() ||
          m_kerasConfigOutputName.empty() ||
          m_strScoreCut.empty() ||
          m_strMassCutLow.empty() ||
          m_strMassCutHigh.empty() ||
          m_decorationName.empty() ||
          m_weightFileName.empty()) ||
        ((m_weightDecorationName.empty() ||
          m_weightHistogramName.empty() ||
          m_weightFlavors.empty()) && m_calcSF)
       )
    {
      ATH_MSG_ERROR( "No config file provided OR you haven't manually specified all needed parameters" ) ;
      ATH_MSG_ERROR( "Please read the TWiki for this tool" );
      return StatusCode::FAILURE;
    }
 
  }
 
  ATH_MSG_INFO( "Mass cut low   : " << m_strMassCutLow );
  ATH_MSG_INFO( "Mass cut High  : " << m_strMassCutHigh );
  ATH_MSG_INFO( "Score cut low  : " << m_strScoreCut );
 
  if ( m_kerasCalibArea.empty() ) {
    ATH_MSG_ERROR( "You need to specify where the calibarea is as either being Local or on CVMFS" );
    return StatusCode::FAILURE;
  }
  else if ( !m_kerasCalibArea.compare("Local") ){
    std::string localCalibArea = "BoostedJetTaggers/JSSWTopTaggerDNN/";
    ATH_MSG_INFO( "Using Local calibarea " << localCalibArea );
    m_kerasConfigFilePath = PathResolverFindCalibFile(localCalibArea+m_kerasConfigFileName);
    if ( m_calcSF )
      m_weightConfigPath = PathResolverFindCalibFile(localCalibArea+m_weightFileName);
  }
  else {
    ATH_MSG_INFO( "Using CVMFS calibarea" );
    m_kerasConfigFilePath = PathResolverFindCalibFile( (m_kerasCalibArea+m_kerasConfigFileName).c_str() );
    if ( m_calcSF )
      m_weightConfigPath = PathResolverFindCalibFile( (m_kerasCalibArea+m_weightFileName).c_str());
  }
 
  ATH_MSG_INFO( "DNN Tagger configured with: " << m_kerasConfigFilePath );
 
  std::ifstream input_cfg( m_kerasConfigFilePath.c_str() );
 
  if ( !input_cfg.is_open() ) {
    ATH_MSG_ERROR( "Error openning config file: " << m_kerasConfigFilePath );
    ATH_MSG_ERROR( "Are you sure that the file exists at this path?" );
    return StatusCode::FAILURE;
  }
 
  lwt::JSONConfig cfg = lwt::parse_json( input_cfg );
 
  ATH_MSG_INFO( "Keras Network NLayers: " << cfg.layers.size() );
 
  m_lwnn = std::make_unique<lwt::LightweightNeuralNetwork>(cfg.inputs, cfg.layers, cfg.outputs);
 
  if ( !m_tagType.compare("TopQuark") ) {
    ATH_MSG_DEBUG( "This is a top quark tagger" );
    m_tagClass = TAGCLASS::TopQuark;
  }
  else if ( !m_tagType.compare("WBoson") ) {
    ATH_MSG_DEBUG( "This is a W boson tagger" );
    m_tagClass = TAGCLASS::WBoson;
  }
  else if ( !m_tagType.compare("ZBoson") ) {
    ATH_MSG_DEBUG( "This is a Z boson tagger" );
    m_tagClass = TAGCLASS::ZBoson;
  }
  else {
    ATH_MSG_ERROR( "I can't tell what kind of tagger your configuration is for." );
    return StatusCode::FAILURE;
  }
 
  m_acceptInfo.addCut( "PassMassLow" , "mJet > mCutLow" );
  m_acceptInfo.addCut( "PassScore"   , "ScoreJet > ScoreCut" );
  if ( m_tagClass == TAGCLASS::WBoson || m_tagClass == TAGCLASS::ZBoson ) {
    m_acceptInfo.addCut( "PassMassHigh", "mJet < mCutHigh" );
  }
 
  printCuts();
 
  ATH_CHECK( JSSTaggerBase::initialize() );
 
#ifndef XAOD_STANDALONE
  if (m_suppressOutputDependence) {
    renounce(m_decPassMassKey);
    renounce(m_decPassScoreKey);
    renounce(m_decTaggedKey);
    renounce(m_decCutMLowKey);
    renounce(m_decCutMHighKey);
    renounce(m_decScoreCutKey);
    renounce(m_decScoreValueKey);
  }
#endif
 
  ATH_MSG_INFO( "DNN Tagger tool initialized" );
 
  return StatusCode::SUCCESS;
 
}
 
StatusCode JSSWTopTaggerDNN::tag( const xAOD::Jet& jet ) const {
 
  ATH_MSG_DEBUG( "Obtaining DNN result" );
 
  asg::AcceptData acceptData( &m_acceptInfo );
 
  ATH_CHECK( resetCuts( acceptData ) );
 
  ATH_CHECK( checkKinRange( jet, acceptData ) );
 
  float jet_pt   = jet.pt()/1000.0;
  float jet_mass = jet.m()/1000.0;
 
  float jet_score = getScore(jet);
 
  float cut_mass_low  = m_funcMassCutLow ->Eval(jet_pt);
  float cut_mass_high = m_funcMassCutHigh->Eval(jet_pt);
  float cut_score     = m_funcScoreCut   ->Eval(jet_pt);
 
  ATH_MSG_VERBOSE( "Cut values : Mass window = [" << cut_mass_low << "," << cut_mass_high << "], score cut = " << cut_score );
  ATH_MSG_VERBOSE( "Jet values : Mass = " << jet_mass << ", score = " << jet_score );
 
  ATH_CHECK( getWeight( jet, jet_score > cut_score, acceptData ) );
 
  ATH_MSG_DEBUG( "Decorating with score" );
 
  SG::WriteDecorHandle<xAOD::JetContainer, bool> decPassMass(m_decPassMassKey);
  SG::WriteDecorHandle<xAOD::JetContainer, bool> decPassScore(m_decPassScoreKey);
  SG::WriteDecorHandle<xAOD::JetContainer, bool> decTagged(m_decTaggedKey);
  SG::WriteDecorHandle<xAOD::JetContainer, float> decCutMLow(m_decCutMLowKey);
  SG::WriteDecorHandle<xAOD::JetContainer, float> decCutMHigh(m_decCutMHighKey);
  SG::WriteDecorHandle<xAOD::JetContainer, float> decScoreCut(m_decScoreCutKey);
  SG::WriteDecorHandle<xAOD::JetContainer, float> decScoreValue(m_decScoreValueKey);
 
  decCutMLow(jet) = cut_mass_low;
  decCutMHigh(jet) = cut_mass_high;
  decScoreCut(jet) = cut_score;
  decScoreValue(jet) = jet_score;
 
  bool passCuts = true;
 
  if ( m_tagClass == TAGCLASS::WBoson || m_tagClass == TAGCLASS::ZBoson ) {
    ATH_MSG_VERBOSE( "Determining WZ tag return" );
    if ( jet_mass > cut_mass_low ) acceptData.setCutResult( "PassMassLow", true );
    if ( jet_mass < cut_mass_high ) acceptData.setCutResult( "PassMassHigh", true );
    if ( jet_score > cut_score ) acceptData.setCutResult( "PassScore", true );
    decPassMass(jet) = acceptData.getCutResult( "PassMassLow" ) && acceptData.getCutResult( "PassMassHigh" );
    passCuts = passCuts && acceptData.getCutResult( "PassMassLow" ) && acceptData.getCutResult( "PassMassHigh" );
  }
  else if ( m_tagClass == TAGCLASS::TopQuark ) {
    ATH_MSG_VERBOSE( "Determining TopQuark tag return" );
    if ( jet_mass > cut_mass_low ) acceptData.setCutResult( "PassMassLow", true );
    if ( jet_score > cut_score ) acceptData.setCutResult( "PassScore", true );
    decPassMass(jet) = acceptData.getCutResult( "PassMassLow" );
    passCuts = passCuts && acceptData.getCutResult( "PassMassLow" );
  }
 
  decPassScore(jet) = acceptData.getCutResult( "PassScore" );
 
  passCuts = passCuts && acceptData.getCutResult( "PassScore" );
 
  decTagged(jet) = passCuts;
 
  return StatusCode::SUCCESS;
 
}
 
double JSSWTopTaggerDNN::getScore( const xAOD::Jet& jet ) const {
 
  std::map<std::string,double> DNN_inputValues = getJetProperties(jet);
 
  lwt::ValueMap discriminant = m_lwnn->compute(DNN_inputValues);
 
  double DNNscore = -666.;
 
  bool validVars = true;
  SG::ReadDecorHandle<xAOD::JetContainer, float> readTau21WTA(m_readTau21WTAKey);
  if ( readTau21WTA(jet) < 0.0 ) validVars = false;
  if ( m_tagClass == TAGCLASS::TopQuark ) {
    SG::ReadDecorHandle<xAOD::JetContainer, float> readTau32WTA(m_readTau32WTAKey);
    if ( readTau32WTA(jet) < 0.0 ) validVars = false;
  }
 
  if ( !validVars ) {
 
    if ( m_nWarnVar++ < m_nWarnMax ) ATH_MSG_WARNING( "One (or more) tagger input variable has an out-of-range value, setting score to -666" );
    else ATH_MSG_DEBUG( "One (or more) tagger input variable has an out-of-range value, setting score to -666" );
 
    return DNNscore;
 
  }
 
  DNNscore = discriminant[m_kerasConfigOutputName];
 
  return DNNscore;
 
}
 
std::map<std::string,double> JSSWTopTaggerDNN::getJetProperties( const xAOD::Jet& jet ) const {
 
  std::map<std::string,double> DNN_inputValues;
 
  calculateJSSRatios(jet);
 
  ATH_MSG_DEBUG( "Loading variables for common DNN tagger" );
 
  SG::ReadDecorHandle<xAOD::JetContainer, float> readTau21WTA(m_readTau21WTAKey);
  SG::ReadDecorHandle<xAOD::JetContainer, float> readC2(m_readC2Key);
  SG::ReadDecorHandle<xAOD::JetContainer, float> readD2(m_readD2Key);
  SG::ReadDecorHandle<xAOD::JetContainer, float> readSplit12(m_readSplit12Key);
 
  DNN_inputValues["CaloTACombinedMassUncorrelated"] = jet.m();
  DNN_inputValues["JetpTCorrByCombinedMass"] = jet.pt();
 
  DNN_inputValues["Split12"] = readSplit12(jet);
 
  DNN_inputValues["C2"] = readC2(jet);
  DNN_inputValues["D2"] = readD2(jet);
 
  DNN_inputValues["Tau21_wta"] = readTau21WTA(jet);
 
  if ( m_tagClass == TAGCLASS::WBoson ) {
 
    ATH_MSG_DEBUG( "Loading variables for W boson tagger" );
 
    DNN_inputValues["FoxWolfram20"] = jet.getAttribute<float>("FoxWolfram2") / jet.getAttribute<float>("FoxWolfram0");
    DNN_inputValues["PlanarFlow"] = jet.getAttribute<float>("PlanarFlow");
    DNN_inputValues["Angularity"] = jet.getAttribute<float>("Angularity");
    DNN_inputValues["Aplanarity"] = jet.getAttribute<float>("Aplanarity");
    DNN_inputValues["ZCut12"] = jet.getAttribute<float>("ZCut12");
    DNN_inputValues["KtDR"] = jet.getAttribute<float>("KtDR");
 
    int pv_location = findPV();
 
    if(pv_location != -1){
      if( GetUnGroomTracks(jet, pv_location).isSuccess()){
        SG::ReadDecorHandle<xAOD::JetContainer, int> readNtrk500(m_readNtrk500Key);
 
        DNN_inputValues["Ntrk500"] = readNtrk500(jet);
      }
      else{
        ATH_MSG_ERROR("Either the ungroomed parent jet doesn't have 'NumTrkPt500' as an attribute or the parent link is broken");
        DNN_inputValues["Ntrk500"] = -999;
      }
    }
    else {
      ATH_MSG_ERROR("Could not find a primary vertex");
      DNN_inputValues["Ntrk500"] = -999;
    }
 
 
  }
 
  else if ( m_tagClass == TAGCLASS::TopQuark ) {
 
    ATH_MSG_DEBUG("Loading variables for top quark tagger");
 
    SG::ReadDecorHandle<xAOD::JetContainer, float> readTau1WTA(m_readTau1WTAKey);
    SG::ReadDecorHandle<xAOD::JetContainer, float> readTau2WTA(m_readTau2WTAKey);
    SG::ReadDecorHandle<xAOD::JetContainer, float> readTau3WTA(m_readTau3WTAKey);
    SG::ReadDecorHandle<xAOD::JetContainer, float> readTau4WTA(m_readTau4WTAKey);
    SG::ReadDecorHandle<xAOD::JetContainer, float> readTau32WTA(m_readTau32WTAKey);
    SG::ReadDecorHandle<xAOD::JetContainer, float> readTau42WTA(m_readTau42WTAKey);
    SG::ReadDecorHandle<xAOD::JetContainer, float> readSplit23(m_readSplit23Key);
    SG::ReadDecorHandle<xAOD::JetContainer, float> readQw(m_readQwKey);
    SG::ReadDecorHandle<xAOD::JetContainer, float> readThrustMaj(m_readThrustMajKey);
    SG::ReadDecorHandle<xAOD::JetContainer, float> readE3(m_readE3Key);
    SG::ReadDecorHandle<xAOD::JetContainer, float> readL2(m_readL2Key);
    SG::ReadDecorHandle<xAOD::JetContainer, float> readL3(m_readL3Key);
 
 
 
    DNN_inputValues["m"] = jet.m();
    DNN_inputValues["pt"] = jet.pt();
 
    DNN_inputValues["Split23"] = readSplit23(jet);
 
    DNN_inputValues["e3"] = readE3(jet);
 
    DNN_inputValues["Tau1_wta"] = readTau1WTA(jet);
    DNN_inputValues["Tau2_wta"] = readTau2WTA(jet);
    DNN_inputValues["Tau3_wta"] = readTau3WTA(jet);
    if(readTau4WTA.isAvailable()){
      DNN_inputValues["Tau4_wta"] = readTau4WTA(jet);
      DNN_inputValues["Tau42_wta"] = readTau42WTA(jet);
    }
 
    DNN_inputValues["Tau32_wta"] = readTau32WTA(jet);
 
    DNN_inputValues["Qw"] = readQw(jet);
 
    if(readThrustMaj.isAvailable()){
      DNN_inputValues["ThrustMaj"] = readThrustMaj(jet);
    }
    if(readL2.isAvailable()){
      DNN_inputValues["L2"] = readL2(jet);
    }
    if(readL3.isAvailable()){
      DNN_inputValues["L3"] = readL3(jet);
    }
  }
 
  else {
    ATH_MSG_ERROR( "Loading variables failed because the tagger type is not supported" );
  }
 
  return DNN_inputValues;
 
}