TrackTrigger/interface/L1TrackQuality.h

0001 /*
0002 Track Quality Header file
0003
0004 C.Brown 28/07/20
0005 */
0006
0007 #ifndef L1Trigger_TrackTrigger_interface_L1TrackQuality_h
0008 #define L1Trigger_TrackTrigger_interface_L1TrackQuality_h
0009
0010 #include <iostream>
0011 #include <set>
0012 #include <vector>
0013 #include <memory>
0014 #include <string>
0015
0016 #include "FWCore/Framework/interface/Event.h"
0017 #include "FWCore/Framework/interface/EventSetup.h"
0018 #include "FWCore/Framework/interface/Frameworkfwd.h"
0019 #include "FWCore/Framework/interface/MakerMacros.h"
0020 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0021 #include "DataFormats/L1TrackTrigger/interface/TTTrack.h"
0022 #include "DataFormats/L1TrackTrigger/interface/TTTypes.h"
0023 #include "L1Trigger/TrackTrigger/interface/HitPatternHelper.h"
0024 #include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h"
0025 #include <memory>
0026
0027 class L1TrackQuality {
0028 public:
0029   // Enum class used for determining prediction behaviour in setL1TrackQuality
0030   enum class QualityAlgorithm { Cut, GBDT, NN, None };
0031
0032   //Default Constructor
0033   L1TrackQuality();
0034
0035   L1TrackQuality(const edm::ParameterSet& qualityParams);
0036
0037   //Default Destructor
0038   ~L1TrackQuality() = default;
0039
0040   // Controls the conversion between TTTrack features and ML model training features
0041   std::vector<float> featureTransform(TTTrack<Ref_Phase2TrackerDigi_>& aTrack,
0042                                       std::vector<std::string> const& featureNames);
0043
0044   // Passed by reference a track without MVA filled, method fills the track's MVA field
0045   void setL1TrackQuality(TTTrack<Ref_Phase2TrackerDigi_>& aTrack);
0046
0047   // To set private member data
0048   void setCutParameters(std::string const& AlgorithmString,
0049                         float maxZ0,
0050                         float maxEta,
0051                         float chi2dofMax,
0052                         float bendchi2Max,
0053                         float minPt,
0054                         int nStubmin);
0055
0056   void setONNXModel(std::string const& AlgorithmString,
0057                     edm::FileInPath const& ONNXmodel,
0058                     std::string const& ONNXInputName,
0059                     std::vector<std::string> const& featureNames);
0060
0061   void beginRun(const hph::Setup* setup);
0062
0063 private:
0064   // Private Member Data
0065   QualityAlgorithm qualityAlgorithm_ = QualityAlgorithm::None;
0066   edm::FileInPath ONNXmodel_;
0067   std::string ONNXInputName_;
0068   std::vector<std::string> featureNames_;
0069   float maxZ0_;
0070   float maxEta_;
0071   float chi2dofMax_;
0072   float bendchi2Max_;
0073   float minPt_;
0074   int nStubsmin_;
0075   const hph::Setup* setupHPH_;
0076   bool useHPH_;
0077   std::unique_ptr<cms::Ort::ONNXRuntime> runTime_;
0078 };
0079 #endif