GNN.cxx

Go to the documentation of this file.

/*
  Copyright (C) 2002-2025 CERN for the benefit of the ATLAS collaboration
*/
 
#include "FlavorTagInference/GNN.h"
#include "FlavorTagInference/FPCompressionUtils.h"
#include "FlavorTagInference/SaltModel.h"
#include "FlavorTagInference/GNNOptions.h"
#include "FlavorTagInference/StringUtils.h"
#include "xAODJet/JetContainer.h"
 
#include "PathResolver/PathResolver.h"
 
namespace {
  auto getSaltModel(const std::string& nn_file) {
    using namespace FlavorTagInference;
    std::string fullPathToOnnxFile = PathResolverFindCalibFile(nn_file);
    auto saltShared = std::make_shared<const SaltModel>(fullPathToOnnxFile);
    ISaltModelPtr retPtr = saltShared;
    return retPtr;
  }
 
  template <typename T>
  std::string join(const T& c) {
    std::string out;
    std::string sep = "";
    for (const auto& [k, v]: c) {
      out.append(sep);
      sep = ", ";
      out.append(k);
    }
    return out;
  }
}
 
namespace FlavorTagInference {
 
  GNN::GNN(const std::string& nn_file, const GNNOptions& o):
    GNN(getSaltModel(nn_file), o)
  {
  }
 
  GNN::GNN(const GNN& old, const GNNOptions& o):
    GNN(old.m_saltModel, o)
  {
  }
 
  GNN::GNN(ISaltModelPtr util, const GNNOptions& o):
    m_saltModel(util),
    m_dataLoader(util, o)
  {
    // Retrieve the configuration for the model outputs.
    OutputConfig gnn_output_config = m_saltModel->getOutputConfig();
 
    // Create the output decorators.
    auto [dd, rd] = createDecorators(gnn_output_config, m_dataLoader.ftag_options);
    m_dataLoader.data_dependency_names += dd;
 
    // Update dependencies and used remap from the constituents loaders.
    for (const auto& loader : m_dataLoader.vectorVarLoaders) {
      auto loader_name = loader.first;
      auto constituent_loader = loader.second;
      m_dataLoader.data_dependency_names += constituent_loader->getDependencies();
      std::set<std::string> used_remap = constituent_loader->getUsedRemap();
      rd.merge(used_remap);
    }
    dataprep::checkForUnusedRemaps(m_dataLoader.ftag_options.remap_scalar, rd);
 
    // Build the default decorators. Note that this _must_ be called
    // after createDecorators.
    auto unused_defaults = o.default_output_values;
    for (const auto& [name, dec]: m_decorators.jetFloat) {
      float default_value = o.default_output_value;
      if (auto def = unused_defaults.extract(name)) {
        default_value = def.mapped();
      }
      m_defaultValues.emplace_back(dec, default_value);
    }
    if (!unused_defaults.empty()) {
      throw std::runtime_error(
        "unused default values: [" + join(unused_defaults) + "]");
    }
    m_defaultZeroTracks = o.default_zero_tracks;
  }
 
  GNN::GNN(const std::string& file,
           const FlipTagConfig& flip,
           const std::map<std::string, std::string>& remap,
           float def_out_val):
    GNN( file, GNNOptions { flip, remap, def_out_val, {}, false} )
  {}
 
  GNN::GNN(GNN&&) = default;
  GNN::GNN(const GNN&) = default;
  GNN::~GNN() = default;
 
  void GNN::decorateWithDefaults(const xAOD::IParticle& i_jet) const {
    for (const auto& [dec, v]: m_defaultValues) {
      dec(i_jet) = v;
    }
    // for some networks we need to set a lot of empty vectors as well
    if (m_saltModel->getSaltModelVersion() == SaltModelVersion::V1
     || m_saltModel->getSaltModelVersion() == SaltModelVersion::V2) {
      // vector outputs, e.g. track predictions
      for (const auto& dec: m_decorators.jetVecChar) {
        dec.second(i_jet) = {};
      }
      for (const auto& dec: m_decorators.jetVecFloat) {
        dec.second(i_jet) = {};
      }
      for (const auto& [name, loader] : m_dataLoader.vectorVarLoaders) {
        loader->setDefaults(i_jet);
      }
      for (const auto& dec: m_decorators.jetVecInt) {
        dec.second(i_jet) = {};
      }
      for (const auto& dec: m_decorators.jetVecTruncFloat) {
        dec.second(i_jet) = {};
      }
    }
  }
 
  void GNN::decorate(const xAOD::IParticle& i_jet) const {
    /* Main function for decorating a i_jet object with GNN outputs. */
    SaltModelData salt_model_data = m_dataLoader.loadInputs(&i_jet);
    // DumpGnnInputs(salt_model_data.gnn_inputs);
 
    // run inference
    // -------------
    if (m_defaultZeroTracks && salt_model_data.num_inputs == 0) {
      this->decorateWithDefaults(i_jet);
      return;
    }
    auto [out_f, out_vc, out_vf] = m_saltModel->runInference(salt_model_data.gnn_inputs);
 
    // decorate outputs
    // ----------------
 
    // with old metadata, doesn't support writing aux tasks
    if (m_saltModel->getSaltModelVersion() == SaltModelVersion::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(i_jet) = out_vf.at(dec.first).at(0);
      }
    }
    // the new metadata format supports writing aux tasks
    else if (m_saltModel->getSaltModelVersion() == SaltModelVersion::V1
          || m_saltModel->getSaltModelVersion() == SaltModelVersion::V2) {
      // float outputs, e.g. i_jet probabilities
      for (const auto& dec: m_decorators.jetFloat) {
        dec.second(i_jet) = out_f.at(dec.first);
      }
      // vector outputs, e.g. track predictions
      for (const auto& dec: m_decorators.jetVecChar) {
        dec.second(i_jet) = out_vc.at(dec.first);
      }
      for (const auto& dec: m_decorators.jetVecFloat) {
        dec.second(i_jet) = out_vf.at(dec.first);
      }
      for (const auto& dec: m_decorators.jetVecInt) {
        const auto& floats = out_vf.at(dec.first);
        std::vector<int> ints(floats.begin(), floats.end());
        dec.second(i_jet) = ints;
      }
      for (size_t idx = 0; idx < m_decorators.jetVecTruncFloat.size(); ++idx) {
        const auto& dec = m_decorators.jetVecTruncFloat[idx];
        const auto& floats = out_vf.at(dec.first);
        const float scale = m_truncFloatScales[idx];
        const int E = m_truncFloatBits[idx].first;
        const int M = m_truncFloatBits[idx].second;
        std::vector<float> truncated(floats.size());
        for (size_t i = 0; i < floats.size(); ++i) {
          truncated[i] = FPCompressionUtils::truncateToFloat(floats[i] * scale, E, M);
        }
        dec.second(i_jet) = truncated;
      }
 
    }
    else {
      throw std::logic_error("unsupported ONNX metadata version");
    }
  } // end of decorate()
 
  // Dependencies
  FTagDataDependencyNames GNN::getDependencies() const {
    return m_dataLoader.data_dependency_names;
  }
 
  std::tuple<FTagDataDependencyNames, std::set<std::string>>
  GNN::createDecorators(const OutputConfig& outConfig, const FTagOptions& options) {
    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 SaltModelOutput::OutputType::FLOAT:
          m_decorators.jetFloat.emplace_back(outNode.name, Dec<float>(dec_name));
          break;
        case SaltModelOutput::OutputType::VECCHAR:
          m_decorators.jetVecChar.emplace_back(outNode.name, Dec<std::vector<char>>(dec_name));
          break;
        case SaltModelOutput::OutputType::VECFLOAT:
          m_decorators.jetVecFloat.emplace_back(outNode.name, Dec<std::vector<float>>(dec_name));
          break;
        case SaltModelOutput::OutputType::VECINT:
          m_decorators.jetVecInt.emplace_back(outNode.name, Dec<std::vector<int>>(dec_name));
          break;
        case SaltModelOutput::OutputType::VECTRUNCFLOAT:
          m_decorators.jetVecTruncFloat.emplace_back(outNode.name, Dec<std::vector<float>>(dec_name));
          m_truncFloatScales.push_back(outNode.scale);
          m_truncFloatBits.emplace_back(outNode.exp_bits, outNode.man_bits);
          break;
        default:
          throw std::logic_error("Unknown output data type");
      }
    }
 
    return std::make_tuple(deps, usedRemap);
  }
 
} // end of namespace FlavorTagInference