FinalTrackSelectors/plugins/LwtnnESProducer.cc

0001 #include "FWCore/Framework/interface/ModuleFactory.h"
0002 #include "FWCore/Framework/interface/ESProducer.h"
0003 #include "FWCore/ParameterSet/interface/FileInPath.h"
0004
0005 #include "TrackingTools/Records/interface/TrackingComponentsRecord.h"
0006
0007 //from lwtnn
0008 #include "lwtnn/LightweightNeuralNetwork.hh"
0009 #include "lwtnn/parse_json.hh"
0010 #include <fstream>
0011
0012 class LwtnnESProducer : public edm::ESProducer {
0013 public:
0014   LwtnnESProducer(const edm::ParameterSet& iConfig);
0015   ~LwtnnESProducer() override = default;
0016
0017   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0018
0019   // TODO: Use of TrackingComponentsRecord is as inadequate as the
0020   // placement of this ESProducer in RecoTracker/FinalTrackSelectors
0021   // (but it works, I tried to create a new record but for some reason
0022   // did not get it to work). Especially if this producer gets used
0023   // wider we should figure out a better record and package.
0024   std::unique_ptr<lwt::LightweightNeuralNetwork> produce(const TrackingComponentsRecord& iRecord);
0025
0026 private:
0027   edm::FileInPath fileName_;
0028 };
0029
0030 LwtnnESProducer::LwtnnESProducer(const edm::ParameterSet& iConfig)
0031     : fileName_(iConfig.getParameter<edm::FileInPath>("fileName")) {
0032   auto componentName = iConfig.getParameter<std::string>("ComponentName");
0033   setWhatProduced(this, componentName);
0034 }
0035
0036 void LwtnnESProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0037   edm::ParameterSetDescription desc;
0038   desc.add<std::string>("ComponentName", "lwtnnESProducer");
0039   desc.add<edm::FileInPath>("fileName", edm::FileInPath());
0040   descriptions.add("lwtnnESProducer", desc);
0041 }
0042
0043 std::unique_ptr<lwt::LightweightNeuralNetwork> LwtnnESProducer::produce(const TrackingComponentsRecord& iRecord) {
0044   std::ifstream jsonfile(fileName_.fullPath().c_str());
0045   auto config = lwt::parse_json(jsonfile);
0046   return std::make_unique<lwt::LightweightNeuralNetwork>(config.inputs, config.layers, config.outputs);
0047 }
0048
0049 #include "FWCore/PluginManager/interface/ModuleDef.h"
0050 #include "FWCore/Framework/interface/MakerMacros.h"
0051 DEFINE_FWK_EVENTSETUP_MODULE(LwtnnESProducer);