TICL/interface/TracksterLinkingAlgoBase.h

0001 // Author: Felice Pantaleo - felice.pantaleo@cern.ch
0002 // Date: 09/2018
0003
0004 #ifndef RecoHGCal_TICL_PatternRecognitionAlgoBase_H__
0005 #define RecoHGCal_TICL_PatternRecognitionAlgoBase_H__
0006
0007 #include <memory>
0008 #include <vector>
0009 #include "DataFormats/CaloRecHit/interface/CaloCluster.h"
0010 #include "DataFormats/HGCalReco/interface/Trackster.h"
0011 #include "DataFormats/HGCalReco/interface/TICLLayerTile.h"
0012 #include "DataFormats/HGCalReco/interface/TICLSeedingRegion.h"
0013 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0014 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0015 #include "FWCore/Framework/interface/Event.h"
0016 #include "FWCore/Framework/interface/EventSetup.h"
0017 #include "DataFormats/Common/interface/ValueMap.h"
0018 #include "RecoHGCal/TICL/interface/GlobalCache.h"
0019 #include "DataFormats/HGCalReco/interface/Common.h"
0020 #include "FWCore/Framework/interface/ConsumesCollector.h"
0021 #include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
0022 #include "CommonTools/RecoAlgos/interface/MultiVectorManager.h"
0023 #include "MagneticField/Engine/interface/MagneticField.h"
0024 #include "TrackingTools/GeomPropagators/interface/Propagator.h"
0025 #include "RecoLocalCalo/HGCalRecAlgos/interface/RecHitTools.h"
0026 #include "Geometry/HGCalCommonData/interface/HGCalDDDConstants.h"
0027
0028 namespace edm {
0029   class Event;
0030   class EventSetup;
0031 }  // namespace edm
0032
0033 namespace cms {
0034   namespace Ort {
0035     class ONNXRuntime;
0036   }
0037 }  // namespace cms
0038
0039 namespace ticl {
0040   class TracksterLinkingAlgoBase {
0041   public:
0042     /** \param conf the configuration of the plugin
0043      * \param onnxRuntime the ONNXRuntime, if onnxModelPath was provided in plugin configuration (nullptr otherwise)
0044     */
0045     TracksterLinkingAlgoBase(const edm::ParameterSet& conf,
0046                              edm::ConsumesCollector,
0047                              cms::Ort::ONNXRuntime const* onnxRuntime = nullptr)
0048         : algo_verbosity_(conf.getParameter<int>("algo_verbosity")), onnxRuntime_(onnxRuntime) {}
0049     virtual ~TracksterLinkingAlgoBase() {}
0050
0051     struct Inputs {
0052       const edm::Event& ev;
0053       const edm::EventSetup& es;
0054       const std::vector<reco::CaloCluster>& layerClusters;
0055       const edm::ValueMap<std::pair<float, float>>& layerClustersTime;
0056       const MultiVectorManager<Trackster>& tracksters;
0057
0058       Inputs(const edm::Event& eV,
0059              const edm::EventSetup& eS,
0060              const std::vector<reco::CaloCluster>& lC,
0061              const edm::ValueMap<std::pair<float, float>>& lT,
0062              const MultiVectorManager<Trackster>& tS)
0063           : ev(eV), es(eS), layerClusters(lC), layerClustersTime(lT), tracksters(tS) {}
0064     };
0065
0066     virtual void linkTracksters(const Inputs& input,
0067                                 std::vector<Trackster>& resultTracksters,
0068                                 std::vector<std::vector<unsigned int>>& linkedResultTracksters,
0069                                 std::vector<std::vector<unsigned int>>& linkedTracksterIdToInputTracksterId) = 0;
0070
0071     virtual void initialize(const HGCalDDDConstants* hgcons,
0072                             const hgcal::RecHitTools rhtools,
0073                             const edm::ESHandle<MagneticField> bfieldH,
0074                             const edm::ESHandle<Propagator> propH) = 0;
0075
0076     // To be called by TracksterLinksProducer at the start of TracksterLinksProducer::produce. Subclasses can use this to store Event and EventSetup
0077     virtual void setEvent(edm::Event& iEvent, edm::EventSetup const& iEventSetup) {}
0078
0079     static void fillPSetDescription(edm::ParameterSetDescription& desc) { desc.add<int>("algo_verbosity", 0); };
0080
0081   protected:
0082     int algo_verbosity_;
0083     cms::Ort::ONNXRuntime const* onnxRuntime_;
0084   };
0085 }  // namespace ticl
0086
0087 #endif