JSSMLTool.h

Go to the documentation of this file.

// Dear emacs, this is -*- c++ -*-
// Copyright (C) 2002-2025 CERN for the benefit of the ATLAS collaboration
#ifndef BOOSTEDJETTAGGERS_JSSMLTOOL_H
#define BOOSTEDJETTAGGERS_JSSMLTOOL_H
 
#include "IJSSMLTool.h"
#include "AsgTools/AsgTool.h"
 
// ONNX Runtime include(s).
#include <onnxruntime_cxx_api.h>
 
// xAOD
#include "xAODJet/JetContainer.h"
#include "xAODPFlow/TrackCaloClusterContainer.h"
// System include(s).
#include <memory> //unique_ptr
#include <string>
#include <map>
#include <vector>
#include <cstdint>
 
class TH2D; 
 
 
namespace AthONNX {


  
class JSSMLTool 
  : public asg::AsgTool, 
    virtual public IJSSMLTool {
  ASG_TOOL_CLASS(JSSMLTool, IJSSMLTool)
 
  public:
    JSSMLTool (const std::string& name);

    virtual StatusCode initialize() override;
    virtual double retrieveConstituentsScore(std::vector<TH2D> Images) const override;
    virtual double retrieveConstituentsScore(std::vector<std::vector<float>> constituents) const override;
    virtual double retrieveConstituentsScore(std::vector<std::vector<float>> constituents, std::vector<std::vector<std::vector<float>>> interactions) const override;
    virtual double retrieveConstituentsScore(std::vector<std::vector<float>> constituents, std::vector<std::vector<std::vector<float>>> interactions, std::vector<std::vector<float>> mask) const override;
    virtual double retrieveHighLevelScore(std::map<std::string, double> JSSVars) const override;
 
    virtual std::vector<float> retrieveConstituentsScoreMultiClass(const std::vector<std::vector<float>>& constituents, const std::vector<std::vector<std::vector<float>>>& interactions, const std::vector<std::vector<float>>& mask) const override;
 
    // basic tool functions
    std::vector<float> ReadJetImagePixels( std::vector<TH2D> Images ) const;
    std::vector<float> ReadJSSInputs(std::map<std::string, double> JSSVars) const;
    std::vector<int> ReadOutputLabels() const;
 
    // extra methods
    StatusCode SetScaler(std::map<std::string, std::vector<double>> scaler) override;

    std::unique_ptr< Ort::Session > m_session;
    std::unique_ptr< Ort::Env > m_env;
 
    std::map<std::string, std::vector<double>> m_scaler;
    std::map<int, std::string> m_JSSInputMap;
 
   private:

      std::string m_modelFileName;
      std::string m_pixelFileName;
      std::string m_labelFileName;
 
     // input node info
     std::vector<int64_t> m_input_node_dims;
     size_t m_num_input_nodes{};
     std::vector<const char*> m_input_node_names;
 
     // output node info
     std::vector<int64_t> m_output_node_dims;
     size_t m_num_output_nodes{};
     std::vector<const char*> m_output_node_names;
 
     // some configs
      int m_nPixelsX{}, m_nPixelsY{}, m_nPixelsZ{};
 
     int m_nvars{};
 
   }; // class JSSMLTool
 
} // namespace AthONNX
 
#endif // BOOSTEDJETTAGGERS_JSSMLTOOL_H