GNNTool.h

Go to the documentation of this file.

/*
  Copyright (C) 2002-2024 CERN for the benefit of the ATLAS collaboration
*/
 
#ifndef GNN_TOOL_H
#define GNN_TOOL_H
 
// Tool includes
#include "AsgTools/AsgTool.h"
#include "AsgServices/ServiceHandle.h"
#include "FlavorTagDiscriminants/INNSharingSvc.h"
#include "FlavorTagDiscriminants/IBTagConditionalDecorator.h"
#include "FlavorTagDiscriminants/IJetTagConditionalDecorator.h"
 
#include "FlavorTagDiscriminants/GNNToolifiers.h"
 
// EDM includes
#include "xAODBTagging/BTaggingFwd.h"
#include "xAODJet/JetFwd.h"
 
#include <memory>
#include <string>
#include <map>
 
namespace FlavorTagDiscriminants {
 
  class GNN;
 
  //
  // Tool to to flavor tag jet/btagging object
  // using GNN based taggers
  class GNNTool : public asg::AsgTool,
                  virtual public IBTagConditionalDecorator,
                  virtual public IJetTagConditionalDecorator
  {
 
    ASG_TOOL_CLASS2(
      GNNTool,
      IBTagConditionalDecorator,
      IJetTagConditionalDecorator)
    public:
      GNNTool(const std::string& name);
      ~GNNTool();
 
      StatusCode initialize() override;
 
      virtual void decorate(const xAOD::BTagging& btag) const override;
      virtual void decorate(const xAOD::Jet& jet) const override;
      virtual void decorateWithDefaults(const SG::AuxElement& jet) const override;
      void decorate(const xAOD::Jet& jet, const SG::AuxElement& decorated) const;
 
      virtual std::set<std::string> getDecoratorKeys() const override;
      virtual std::set<std::string> getAuxInputKeys() const override;
      virtual std::set<std::string> getConstituentAuxInputKeys() const override;
 
    private:
 
    ServiceHandle<INNSharingSvc> m_nnsvc {
      this, "nnSharingService", "", "NN sharing service"};
    std::string m_nn_file;
    GNNToolProperties m_props;
    std::shared_ptr<const GNN> m_gnn;
  };
}
#endif