FlavorTagDiscriminants::GNN Class Reference

#include <GNN.h>

Collaboration diagram for FlavorTagDiscriminants::GNN:

Classes
struct	Decorators

Public Member Functions
	GNN (const std::string &nnFile, const GNNOptions &opts)
	GNN (const GNN &, const GNNOptions &opts)
	GNN (const std::string &nnFile, const FlipTagConfig &flip_config=FlipTagConfig::STANDARD, const std::map< std::string, std::string > &variableRemapping={}, const TrackLinkType trackLinkType=TrackLinkType::TRACK_PARTICLE, float defaultOutputValue=NAN)
	GNN (GNN &&)
	GNN (const GNN &)
virtual	~GNN ()
virtual void	decorate (const xAOD::BTagging &btag) const
virtual void	decorate (const xAOD::Jet &jet) const
virtual void	decorateWithDefaults (const SG::AuxElement &jet) const
void	decorate (const xAOD::Jet &jet, const SG::AuxElement &decorated) const
virtual std::set< std::string >	getDecoratorKeys () const
virtual std::set< std::string >	getAuxInputKeys () const
virtual std::set< std::string >	getConstituentAuxInputKeys () const

Public Attributes
std::shared_ptr< const OnnxUtil >	m_onnxUtil

Private Types
using	TPC = xAOD::TrackParticleContainer
using	TrackLinks = std::vector< ElementLink< TPC > >
template<typename T >
using	Dec = SG::AuxElement::Decorator< T >
template<typename T >
using	Decs = std::vector< std::pair< std::string, Dec< T > >>

Private Member Functions
	GNN (std::shared_ptr< const OnnxUtil >, const GNNOptions &opts)
std::tuple< FTagDataDependencyNames, std::set< std::string > >	createDecorators (const OnnxUtil::OutputConfig &outConfig, const FTagOptions &options)

Private Attributes
SG::AuxElement::ConstAccessor< ElementLink< xAOD::JetContainer > >	m_jetLink
std::string	m_input_node_name
std::vector< internal::VarFromBTag >	m_varsFromBTag
std::vector< internal::VarFromJet >	m_varsFromJet
std::vector< std::shared_ptr< IConstituentsLoader > >	m_constituentsLoaders
Decorators	m_decorators
float	m_defaultValue
FTagDataDependencyNames	m_dataDependencyNames

Detailed Description

Definition at line 39 of file GNN.h.

Member Typedef Documentation

Dec

template<typename T >

using FlavorTagDiscriminants::GNN::Dec = SG::AuxElement::Decorator<T>

private

Definition at line 74 of file GNN.h.

Decs

template<typename T >

using FlavorTagDiscriminants::GNN::Decs = std::vector<std::pair<std::string, Dec<T> >>

private

Definition at line 77 of file GNN.h.

TPC

using FlavorTagDiscriminants::GNN::TPC = xAOD::TrackParticleContainer

private

Definition at line 70 of file GNN.h.

TrackLinks

using FlavorTagDiscriminants::GNN::TrackLinks = std::vector<ElementLink<TPC> >

private

Definition at line 71 of file GNN.h.

Constructor & Destructor Documentation

GNN() [1/6]

FlavorTagDiscriminants::GNN::GNN	(	const std::string &	nnFile,
		const GNNOptions &	opts
	)

Definition at line 28 of file GNN.cxx.

                                                       :
    GNN(getOnnxUtil(nn_file), o)
  {
  }

GNN() [2/6]

FlavorTagDiscriminants::GNN::GNN	(	const GNN &	old,
		const GNNOptions &	opts
	)

Definition at line 33 of file GNN.cxx.

                                             :
    GNN(old.m_onnxUtil, o)
  {
  }

GNN() [3/6]

FlavorTagDiscriminants::GNN::GNN	(	const std::string &	nnFile,
		const FlipTagConfig &	flip_config = FlipTagConfig::STANDARD,
		const std::map< std::string, std::string > &	variableRemapping = {},
		const TrackLinkType	trackLinkType = TrackLinkType::TRACK_PARTICLE,
		float	defaultOutputValue = NAN
	)

Definition at line 82 of file GNN.cxx.

                             :
    GNN( file, GNNOptions { flip, remap, link_type, def_out_val} )
  {}

GNN() [4/6]

FlavorTagDiscriminants::GNN::GNN ( GNN &&  )

default

GNN() [5/6]

FlavorTagDiscriminants::GNN::GNN ( const GNN &  )

default

~GNN()

FlavorTagDiscriminants::GNN::~GNN ( )

virtualdefault

GNN() [6/6]

FlavorTagDiscriminants::GNN::GNN ( std::shared_ptr< const OnnxUtil >  util, const GNNOptions &  opts  )

private

Definition at line 38 of file GNN.cxx.

                                                                 :
    m_onnxUtil(util),
    m_jetLink(jetLinkName),
    m_defaultValue(o.default_output_value)
  {
 
    // Extract metadata from the ONNX file, primarily about the model's inputs.
    auto lwt_config = m_onnxUtil->getLwtConfig();
 
    // Create configuration objects for data preprocessing.
    auto [inputs, constituents_configs, options] = dataprep::createGetterConfig(
        lwt_config, o.flip_config, o.variable_remapping, o.track_link_type);
    for (auto config : constituents_configs){
      switch (config.type){
      case ConstituentsType::TRACK:
        m_constituentsLoaders.push_back(std::make_shared<TracksLoader>(config, options));
        break;
      case ConstituentsType::IPARTICLE:
        m_constituentsLoaders.push_back(std::make_shared<IParticlesLoader>(config, options));
        break;
      }
    }
 
    // Initialize jet and b-tagging input getters.
    auto [vb, vj, ds] = dataprep::createBvarGetters(inputs);
    m_varsFromBTag = vb;
    m_varsFromJet = vj;
    m_dataDependencyNames = ds;
 
    // Retrieve the configuration for the model outputs.
    OnnxUtil::OutputConfig gnn_output_config = m_onnxUtil->getOutputConfig();
 
    // Create the output decorators.
    auto [dd, rd] = createDecorators(gnn_output_config, options);
    m_dataDependencyNames += dd;
 
    // Update dependencies and used remap from the constituents loaders.
    for (auto loader : m_constituentsLoaders){
      m_dataDependencyNames += loader->getDependencies();
      rd.merge(loader->getUsedRemap());
    }
    dataprep::checkForUnusedRemaps(options.remap_scalar, rd);
  }

Member Function Documentation

createDecorators()

std::tuple< FTagDataDependencyNames, std::set< std::string > > FlavorTagDiscriminants::GNN::createDecorators ( const OnnxUtil::OutputConfig &  outConfig, const FTagOptions &  options  )

private

Definition at line 231 of file GNN.cxx.

                                                                                         {
    FTagDataDependencyNames deps;
    Decorators decs;
 
    std::map<std::string, std::string> remap = options.remap_scalar;
    std::set<std::string> usedRemap;
 
    // get the regex to rewrite the outputs if we're using flip taggers
    auto flip_converters = dataprep::getNameFlippers(options.flip);
    std::string context = "building negative tag b-btagger";
 
    for (const auto& outNode : outConfig) {
      // the node's output name will be used to define the decoration name
      std::string dec_name = outNode.name;
 
      // modify the deco name if we're using flip taggers
      if (options.flip != FlipTagConfig::STANDARD) {
        dec_name = str::sub_first(flip_converters, dec_name, context);
      }
 
      // remap the deco name if necessary
      dec_name = str::remapName(dec_name, remap, usedRemap);
 
      // keep track of dependencies for EDM bookkeeping
      deps.bTagOutputs.insert(dec_name);
 
      // Create decorators based on output type and target
      switch (outNode.type) {
        case OnnxOutput::OutputType::FLOAT:
          m_decorators.jetFloat.emplace_back(outNode.name, Dec<float>(dec_name));
          break;
        case OnnxOutput::OutputType::VECCHAR:
          m_decorators.jetVecChar.emplace_back(outNode.name, Dec<std::vector<char>>(dec_name));
          break;
        case OnnxOutput::OutputType::VECFLOAT:
          m_decorators.jetVecFloat.emplace_back(outNode.name, Dec<std::vector<float>>(dec_name));
          break;
        default:
          throw std::logic_error("Unknown output data type");
      }
    }
 
    // Create decorators for links to the input tracks
    if (!m_decorators.jetVecChar.empty() || !m_decorators.jetVecFloat.empty()) {
      std::string name = m_onnxUtil->getModelName() + "_TrackLinks";
      name = str::remapName(name, remap, usedRemap);
      deps.bTagOutputs.insert(name);
      m_decorators.jetTrackLinks.emplace_back(name, Dec<TrackLinks>(name));
    }
 
    return std::make_tuple(deps, usedRemap);
  }

decorate() [1/3]

void FlavorTagDiscriminants::GNN::decorate ( const xAOD::BTagging &  btag ) const

virtual

Definition at line 94 of file GNN.cxx.

                                                   {
    /* tag a b-tagging object */
    auto jetLink = m_jetLink(btag);
    if (!jetLink.isValid()) {
      throw std::runtime_error("invalid jetLink");
    }
    const xAOD::Jet& jet = **jetLink;
    decorate(jet, btag);
  }

decorate() [2/3]

void FlavorTagDiscriminants::GNN::decorate ( const xAOD::Jet &  jet ) const

virtual

Definition at line 104 of file GNN.cxx.

                                             {
    /* tag a jet */
    decorate(jet, jet);
  }

decorate() [3/3]

void FlavorTagDiscriminants::GNN::decorate	(	const xAOD::Jet &	jet,
		const SG::AuxElement &	decorated
	)		const

Definition at line 128 of file GNN.cxx.

                                                                       {
    /* Main function for decorating a jet or b-tagging object with GNN outputs. */
    using namespace internal;
 
    // prepare input
    // -------------
    std::map<std::string, Inputs> gnn_inputs;
 
    // jet level inputs
    std::vector<float> jet_feat;
    for (const auto& getter: m_varsFromBTag) {
      jet_feat.push_back(getter(btag).second);
    }
    for (const auto& getter: m_varsFromJet) {
      jet_feat.push_back(getter(jet).second);
    }
    std::vector<int64_t> jet_feat_dim = {1, static_cast<int64_t>(jet_feat.size())};
    Inputs jet_info(jet_feat, jet_feat_dim);
    if (m_onnxUtil->getOnnxModelVersion() == OnnxModelVersion::V2) {
      gnn_inputs.insert({"jets", jet_info});
    } else {
      gnn_inputs.insert({"jet_features", jet_info});
    }
 
    // constituent level inputs
    Tracks input_tracks;
    for (auto loader : m_constituentsLoaders){
      auto [input_name, input_data, input_objects] = loader->getData(jet, btag);
      if (m_onnxUtil->getOnnxModelVersion() != OnnxModelVersion::V2) {
        input_name.pop_back();
        input_name.append("_features");
      }
      gnn_inputs.insert({input_name, input_data});
      
      // for now we only collect tracks for aux task decoration
      // they have to be converted back from IParticle to TrackParticle first
      if (loader->getType() == ConstituentsType::TRACK){
        for (auto constituent : input_objects){
          input_tracks.push_back(dynamic_cast<const xAOD::TrackParticle*>(constituent));
        }
      }
    }
 
    // run inference
    // -------------
    auto [out_f, out_vc, out_vf] = m_onnxUtil->runInference(gnn_inputs);
 
    // decorate outputs
    // ----------------
 
    // with old metadata, doesn't support writing aux tasks
    if (m_onnxUtil->getOnnxModelVersion() == OnnxModelVersion::V0) {
      for (const auto& dec: m_decorators.jetFloat) {
        if (out_vf.at(dec.first).size() != 1){
          throw std::logic_error("expected vectors of length 1 for float decorators");
        }
        dec.second(btag) = out_vf.at(dec.first).at(0);
      }
    }
    // the new metadata format supports writing aux tasks
    else if (m_onnxUtil->getOnnxModelVersion() == OnnxModelVersion::V1) {
      // float outputs, e.g. jet probabilities
      for (const auto& dec: m_decorators.jetFloat) {
        dec.second(btag) = out_f.at(dec.first);
      }
      // vector outputs, e.g. track predictions
      for (const auto& dec: m_decorators.jetVecChar) {
        dec.second(btag) = out_vc.at(dec.first);
      }
      for (const auto& dec: m_decorators.jetVecFloat) {
        dec.second(btag) = out_vf.at(dec.first);
      }
 
      // decorate links to the input tracks to the b-tagging object
      for (const auto& dec: m_decorators.jetTrackLinks) {
        TrackLinks links;
        for (const xAOD::TrackParticle* it: input_tracks) {
          TrackLinks::value_type link;
          const auto* itc = dynamic_cast<const xAOD::TrackParticleContainer*>(
            it->container());
          link.toIndexedElement(*itc, it->index());
          links.push_back(link);
        }
        dec.second(btag) = links;
      }
    }
    else {
      throw std::logic_error("unsupported ONNX metadata version");
    }
  } // end of decorate()

decorateWithDefaults()

void FlavorTagDiscriminants::GNN::decorateWithDefaults ( const SG::AuxElement &  jet ) const

virtual

Definition at line 109 of file GNN.cxx.

                                                              {
    for (const auto& dec: m_decorators.jetFloat) {
      dec.second(jet) = m_defaultValue;
    }
    // for some networks we need to set a lot of empty vectors as well
    if (m_onnxUtil->getOnnxModelVersion() == OnnxModelVersion::V1) {
      // vector outputs, e.g. track predictions
      for (const auto& dec: m_decorators.jetVecChar) {
        dec.second(jet) = {};
      }
      for (const auto& dec: m_decorators.jetVecFloat) {
        dec.second(jet) = {};
      }
      for (const auto& dec: m_decorators.jetTrackLinks) {
        dec.second(jet) = {};
      }
    }
  }

getAuxInputKeys()

std::set< std::string > FlavorTagDiscriminants::GNN::getAuxInputKeys ( ) const

virtual

Definition at line 223 of file GNN.cxx.

                                               {
    return m_dataDependencyNames.bTagInputs;
  }

getConstituentAuxInputKeys()

std::set< std::string > FlavorTagDiscriminants::GNN::getConstituentAuxInputKeys ( ) const

virtual

Definition at line 226 of file GNN.cxx.

                                                          {
    return m_dataDependencyNames.trackInputs;
  }

getDecoratorKeys()

std::set< std::string > FlavorTagDiscriminants::GNN::getDecoratorKeys ( ) const

virtual

Definition at line 220 of file GNN.cxx.

                                                {
    return m_dataDependencyNames.bTagOutputs;
  }

Member Data Documentation

m_constituentsLoaders

std::vector<std::shared_ptr<IConstituentsLoader> > FlavorTagDiscriminants::GNN::m_constituentsLoaders

private

Definition at line 96 of file GNN.h.

m_dataDependencyNames

FTagDataDependencyNames FlavorTagDiscriminants::GNN::m_dataDependencyNames

private

Definition at line 100 of file GNN.h.

m_decorators

Decorators FlavorTagDiscriminants::GNN::m_decorators

private

Definition at line 98 of file GNN.h.

m_defaultValue

float FlavorTagDiscriminants::GNN::m_defaultValue

private

Definition at line 99 of file GNN.h.

m_input_node_name

std::string FlavorTagDiscriminants::GNN::m_input_node_name

private

Definition at line 93 of file GNN.h.

m_jetLink

SG::AuxElement::ConstAccessor<ElementLink<xAOD::JetContainer> > FlavorTagDiscriminants::GNN::m_jetLink

private

Definition at line 92 of file GNN.h.

m_onnxUtil

std::shared_ptr<const OnnxUtil> FlavorTagDiscriminants::GNN::m_onnxUtil

Definition at line 65 of file GNN.h.

m_varsFromBTag

std::vector<internal::VarFromBTag> FlavorTagDiscriminants::GNN::m_varsFromBTag

private

Definition at line 94 of file GNN.h.

m_varsFromJet

std::vector<internal::VarFromJet> FlavorTagDiscriminants::GNN::m_varsFromJet

private

Definition at line 95 of file GNN.h.

---

The documentation for this class was generated from the following files:

• GNN.h
• GNN.cxx