PatAlgos/plugins/PNETLeptonProducer.cc

0001 // -*- C++ -*-
0002 //
0003 // Package:    PhysicsTools/PatAlgos
0004 // Class:      PNETLeptonProducer
0005 //
0006 // Original Author:  Sergio Sanchez Cruz
0007 //         Created:  Mon, 15 May 2023 08:32:03 GMT
0008 //
0009
0010 #include "FWCore/Framework/interface/Frameworkfwd.h"
0011 #include "FWCore/Framework/interface/stream/EDProducer.h"
0012
0013 #include "FWCore/Framework/interface/Event.h"
0014 #include "FWCore/Framework/interface/MakerMacros.h"
0015
0016 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0017 #include "FWCore/Utilities/interface/StreamID.h"
0018
0019 #include "DataFormats/BTauReco/interface/DeepBoostedJetFeatures.h"
0020 #include "DataFormats/Common/interface/ValueMap.h"
0021 #include "DataFormats/PatCandidates/interface/Electron.h"
0022 #include "DataFormats/PatCandidates/interface/Muon.h"
0023
0024 #include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h"
0025 #include "RecoBTag/FeatureTools/interface/deep_helpers.h"
0026
0027 using namespace cms::Ort;
0028 using namespace btagbtvdeep;
0029
0030 template <typename LeptonType>
0031 class PNETLeptonProducer : public edm::stream::EDProducer<edm::GlobalCache<cms::Ort::ONNXRuntime>> {
0032 public:
0033   explicit PNETLeptonProducer(const edm::ParameterSet &, const cms::Ort::ONNXRuntime *);
0034   ~PNETLeptonProducer() override {}
0035
0036   static void fillDescriptions(edm::ConfigurationDescriptions &descriptions);
0037
0038   static std::unique_ptr<ONNXRuntime> initializeGlobalCache(const edm::ParameterSet &);
0039   static void globalEndJob(const ONNXRuntime *);
0040
0041 private:
0042   using LeptonTagInfoCollection = DeepBoostedJetFeaturesCollection;
0043
0044   void produce(edm::Event &, const edm::EventSetup &) override;
0045   void make_inputs(const DeepBoostedJetFeatures &);
0046
0047   edm::EDGetTokenT<LeptonTagInfoCollection> src_;
0048   edm::EDGetTokenT<edm::View<LeptonType>> leps_;
0049   std::vector<std::string> flav_names_;
0050   std::vector<std::string> input_names_;            // names of each input group - the ordering is important!
0051   std::vector<std::vector<int64_t>> input_shapes_;  // shapes of each input group (-1 for dynamic axis)
0052   std::vector<unsigned> input_sizes_;               // total length of each input vector
0053   std::unordered_map<std::string, btagbtvdeep::PreprocessParams>
0054       prep_info_map_;  // preprocessing info for each input group
0055
0056   cms::Ort::FloatArrays data_;
0057   bool debug_ = false;
0058 };
0059
0060 template <typename LeptonType>
0061 PNETLeptonProducer<LeptonType>::PNETLeptonProducer(const edm::ParameterSet &iConfig, const cms::Ort::ONNXRuntime *cache)
0062     : src_(consumes<LeptonTagInfoCollection>(iConfig.getParameter<edm::InputTag>("src"))),
0063       leps_(consumes<edm::View<LeptonType>>(iConfig.getParameter<edm::InputTag>("srcLeps"))),
0064       flav_names_(iConfig.getParameter<std::vector<std::string>>("flav_names")),
0065       debug_(iConfig.getUntrackedParameter<bool>("debugMode", false)) {
0066   ParticleNetConstructor(iConfig, true, input_names_, prep_info_map_, input_shapes_, input_sizes_, &data_);
0067
0068   if (debug_) {
0069     for (unsigned i = 0; i < input_names_.size(); ++i) {
0070       const auto &group_name = input_names_.at(i);
0071       std::cout << group_name << std::endl;
0072       if (!input_shapes_.empty()) {
0073         std::cout << group_name << "\nshapes: ";
0074         for (const auto &x : input_shapes_.at(i)) {
0075           std::cout << x << ", ";
0076         }
0077       }
0078       std::cout << "\nvariables: ";
0079       for (const auto &x : prep_info_map_.at(group_name).var_names) {
0080         std::cout << x << ", ";
0081       }
0082       std::cout << "\n";
0083     }
0084     std::cout << "\n";
0085   }
0086
0087   for (const auto &flav_name : flav_names_) {
0088     produces<edm::ValueMap<float>>(flav_name);
0089   }
0090 }
0091
0092 template <typename LeptonType>
0093 void PNETLeptonProducer<LeptonType>::produce(edm::Event &iEvent, const edm::EventSetup &iSetup) {
0094   edm::Handle<LeptonTagInfoCollection> src;
0095   iEvent.getByToken(src_, src);
0096
0097   edm::Handle<edm::View<LeptonType>> leps;
0098   iEvent.getByToken(leps_, leps);
0099
0100   std::vector<std::vector<float>> mvaScores(flav_names_.size(), std::vector<float>(leps->size(), -1));
0101
0102   // tagInfo src could be empty if the event has no PV
0103   if (!src->empty()) {
0104     assert(src->size() == leps->size());
0105     for (size_t ilep = 0; ilep < src->size(); ilep++) {
0106       const auto &taginfo = (*src)[ilep];
0107       make_inputs(taginfo);
0108       auto outputs = globalCache()->run(input_names_, data_, input_shapes_)[0];
0109       // std::cout<<"outputs.size(): "<<outputs.size()<<std::endl;
0110       // std::cout<<"flav_names_.size(): "<<flav_names_.size()<<std::endl;
0111       assert(outputs.size() == flav_names_.size());
0112       for (unsigned int iflav = 0; iflav < flav_names_.size(); ++iflav) {
0113         mvaScores[iflav][ilep] = outputs.at(iflav);
0114       }
0115     }
0116   }
0117
0118   for (unsigned int iflav = 0; iflav < flav_names_.size(); ++iflav) {
0119     std::unique_ptr<edm::ValueMap<float>> pnScore(new edm::ValueMap<float>());
0120     edm::ValueMap<float>::Filler filler(*pnScore);
0121     filler.insert(leps, mvaScores[iflav].begin(), mvaScores[iflav].end());
0122     filler.fill();
0123     iEvent.put(std::move(pnScore), flav_names_[iflav]);
0124   }
0125 }
0126
0127 template <typename LeptonType>
0128 std::unique_ptr<cms::Ort::ONNXRuntime> PNETLeptonProducer<LeptonType>::initializeGlobalCache(
0129     const edm::ParameterSet &cfg) {
0130   return std::make_unique<cms::Ort::ONNXRuntime>(cfg.getParameter<edm::FileInPath>("model_path").fullPath());
0131 }
0132
0133 template <typename LeptonType>
0134 void PNETLeptonProducer<LeptonType>::globalEndJob(const cms::Ort::ONNXRuntime *cache) {}
0135
0136 template <typename LeptonType>
0137 void PNETLeptonProducer<LeptonType>::fillDescriptions(edm::ConfigurationDescriptions &descriptions) {
0138   edm::ParameterSetDescription desc;
0139   desc.add<edm::InputTag>("src")->setComment("input variables");
0140   desc.add<edm::InputTag>("srcLeps")->setComment("input physics object collection. src and srcLeps must be in synch");
0141   desc.add<std::vector<std::string>>("flav_names")->setComment("Names of the oputput classes");
0142   desc.add<std::string>("preprocess_json", "PhysicsTools/NanoAOD/data/PNetMuonId/preprocess.json");
0143   desc.add<edm::FileInPath>("model_path", edm::FileInPath("PhysicsTools/NanoAOD/data/PNetMuonId/model.onnx"));
0144   desc.addOptionalUntracked<bool>("debugMode", false);
0145
0146   std::string modname;
0147   if (typeid(LeptonType) == typeid(pat::Muon))
0148     modname += "muon";
0149   else if (typeid(LeptonType) == typeid(pat::Electron))
0150     modname += "electron";
0151   modname += "PNetTags";
0152   descriptions.add(modname, desc);
0153 }
0154
0155 template <typename LeptonType>
0156 void PNETLeptonProducer<LeptonType>::make_inputs(const DeepBoostedJetFeatures &taginfo) {
0157   for (unsigned igroup = 0; igroup < input_names_.size(); ++igroup) {
0158     const auto &group_name = input_names_[igroup];
0159     const auto &prep_params = prep_info_map_.at(group_name);
0160     auto &group_values = data_[igroup];
0161     group_values.resize(input_sizes_[igroup]);
0162
0163     // first reset group_values to 0
0164     std::fill(group_values.begin(), group_values.end(), 0);
0165     unsigned curr_pos = 0;
0166
0167     // transform/pad
0168     for (unsigned i = 0; i < prep_params.var_names.size(); ++i) {
0169       const auto &varname = prep_params.var_names[i];
0170       const auto &raw_value = taginfo.get(varname);
0171       const auto &info = prep_params.info(varname);
0172       int insize = btagbtvdeep::center_norm_pad(raw_value,
0173                                                 info.center,
0174                                                 info.norm_factor,
0175                                                 prep_params.min_length,
0176                                                 prep_params.max_length,
0177                                                 group_values,
0178                                                 curr_pos,
0179                                                 info.pad,
0180                                                 info.replace_inf_value,
0181                                                 info.lower_bound,
0182                                                 info.upper_bound);
0183       curr_pos += insize;
0184       if (i == 0 && (!input_shapes_.empty())) {
0185         input_shapes_[igroup][2] = insize;
0186       }
0187     }
0188     group_values.resize(curr_pos);
0189   }
0190 }
0191
0192 typedef PNETLeptonProducer<pat::Muon> MuonPNETProducer;
0193 typedef PNETLeptonProducer<pat::Electron> ElectronPNETProducer;
0194
0195 //define this as a plug-in
0196 DEFINE_FWK_MODULE(MuonPNETProducer);
0197 DEFINE_FWK_MODULE(ElectronPNETProducer);