PatAlgos/interface/MuonMvaIDEstimator.h

0001 #ifndef PhysicsTools_PatAlgos_MuonMvaIDEstimator_h
0002 #define PhysicsTools_PatAlgos_MuonMvaIDEstimator_h
0003
0004 #include <memory>
0005 #include <string>
0006 #include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h"
0007 #include "FWCore/Framework/interface/stream/EDProducer.h"
0008 #include "FWCore/Utilities/interface/StreamID.h"
0009 #include "FWCore/Framework/interface/stream/EDAnalyzer.h"
0010 //
0011
0012 namespace pat {
0013   class Muon;
0014 }
0015
0016 namespace edm {
0017   class FileInPath;
0018 }
0019
0020 namespace pat {
0021   class MuonMvaIDEstimator {
0022   public:
0023     MuonMvaIDEstimator(const edm::FileInPath &weightsfile);
0024     ~MuonMvaIDEstimator() = default;
0025
0026     static void fillDescriptions(edm::ConfigurationDescriptions &);
0027     static void globalEndJob(const cms::Ort::ONNXRuntime *);
0028     std::vector<float> computeMVAID(const pat::Muon &imuon) const;
0029
0030   private:
0031     std::vector<std::string> flav_names_;   // names of the output scores
0032     std::vector<std::string> input_names_;  // names of each input group - the ordering is important!
0033     std::unique_ptr<const cms::Ort::ONNXRuntime> randomForest_;
0034   };
0035 };  // namespace pat
0036 #endif