ONNXRuntime/plugins/DeepFlavourONNXJetTagsProducer.cc

0001 #include "FWCore/Framework/interface/Frameworkfwd.h"
0002 #include "FWCore/Framework/interface/stream/EDProducer.h"
0003
0004 #include "FWCore/Framework/interface/Event.h"
0005 #include "FWCore/Framework/interface/MakerMacros.h"
0006
0007 #include "FWCore/Framework/interface/makeRefToBaseProdFrom.h"
0008
0009 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0010 #include "FWCore/Utilities/interface/StreamID.h"
0011
0012 #include "DataFormats/BTauReco/interface/JetTag.h"
0013
0014 #include "DataFormats/BTauReco/interface/DeepFlavourTagInfo.h"
0015
0016 #include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h"
0017
0018 using namespace cms::Ort;
0019
0020 class DeepFlavourONNXJetTagsProducer : public edm::stream::EDProducer<edm::GlobalCache<ONNXRuntime>> {
0021 public:
0022   explicit DeepFlavourONNXJetTagsProducer(const edm::ParameterSet&, const ONNXRuntime*);
0023   ~DeepFlavourONNXJetTagsProducer() override;
0024
0025   static void fillDescriptions(edm::ConfigurationDescriptions&);
0026
0027   static std::unique_ptr<ONNXRuntime> initializeGlobalCache(const edm::ParameterSet&);
0028   static void globalEndJob(const ONNXRuntime*);
0029
0030 private:
0031   typedef std::vector<reco::DeepFlavourTagInfo> TagInfoCollection;
0032   typedef reco::JetTagCollection JetTagCollection;
0033
0034   void produce(edm::Event&, const edm::EventSetup&) override;
0035
0036   void make_inputs(unsigned i_jet, const reco::DeepFlavourTagInfo& taginfo);
0037
0038   const edm::EDGetTokenT<TagInfoCollection> src_;
0039   std::vector<std::string> flav_names_;
0040   std::vector<std::string> input_names_;
0041   std::vector<std::string> output_names_;
0042
0043   enum InputIndexes { kGlobal = 0, kChargedCandidates = 1, kNeutralCandidates = 2, kVertices = 3, kJetPt = 4 };
0044   constexpr static unsigned n_features_global_ = 15;
0045   constexpr static unsigned n_cpf_ = 25;
0046   constexpr static unsigned n_features_cpf_ = 16;
0047   constexpr static unsigned n_npf_ = 25;
0048   constexpr static unsigned n_features_npf_ = 6;
0049   constexpr static unsigned n_sv_ = 4;
0050   constexpr static unsigned n_features_sv_ = 12;
0051   constexpr static unsigned n_features_jetpt_ = 1;
0052   const static std::vector<unsigned> input_sizes_;
0053
0054   // hold the input data
0055   FloatArrays data_;
0056 };
0057
0058 const std::vector<unsigned> DeepFlavourONNXJetTagsProducer::input_sizes_{
0059     n_features_global_, n_cpf_* n_features_cpf_, n_npf_* n_features_npf_, n_sv_* n_features_sv_, n_features_jetpt_};
0060
0061 DeepFlavourONNXJetTagsProducer::DeepFlavourONNXJetTagsProducer(const edm::ParameterSet& iConfig,
0062                                                                const ONNXRuntime* cache)
0063     : src_(consumes<TagInfoCollection>(iConfig.getParameter<edm::InputTag>("src"))),
0064       flav_names_(iConfig.getParameter<std::vector<std::string>>("flav_names")),
0065       input_names_(iConfig.getParameter<std::vector<std::string>>("input_names")),
0066       output_names_(iConfig.getParameter<std::vector<std::string>>("output_names")) {
0067   // get output names from flav_names
0068   for (const auto& flav_name : flav_names_) {
0069     produces<JetTagCollection>(flav_name);
0070   }
0071
0072   assert(input_names_.size() == input_sizes_.size());
0073 }
0074
0075 DeepFlavourONNXJetTagsProducer::~DeepFlavourONNXJetTagsProducer() {}
0076
0077 void DeepFlavourONNXJetTagsProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0078   // pfDeepFlavourJetTags
0079   edm::ParameterSetDescription desc;
0080   desc.add<edm::InputTag>("src", edm::InputTag("pfDeepFlavourTagInfos"));
0081   desc.add<std::vector<std::string>>("input_names", {"input_1", "input_2", "input_3", "input_4", "input_5"});
0082   desc.add<edm::FileInPath>(
0083       "model_path", edm::FileInPath("RecoBTag/Combined/data/DeepFlavourV04_12X_training/DeepJet_Run3_122X.onnx"));
0084   desc.add<std::vector<std::string>>("output_names", {"ID_pred/Softmax:0"});
0085   desc.add<std::vector<std::string>>(
0086       "flav_names", std::vector<std::string>{"probb", "probbb", "problepb", "probc", "probuds", "probg"});
0087
0088   descriptions.add("pfDeepFlavourJetTags", desc);
0089 }
0090
0091 std::unique_ptr<ONNXRuntime> DeepFlavourONNXJetTagsProducer::initializeGlobalCache(const edm::ParameterSet& iConfig) {
0092   return std::make_unique<ONNXRuntime>(iConfig.getParameter<edm::FileInPath>("model_path").fullPath());
0093 }
0094
0095 void DeepFlavourONNXJetTagsProducer::globalEndJob(const ONNXRuntime* cache) {}
0096
0097 void DeepFlavourONNXJetTagsProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
0098   edm::Handle<TagInfoCollection> tag_infos;
0099   iEvent.getByToken(src_, tag_infos);
0100
0101   std::vector<std::unique_ptr<JetTagCollection>> output_tags;
0102   if (!tag_infos->empty()) {
0103     // initialize output collection
0104     auto jet_ref = tag_infos->begin()->jet();
0105     auto ref2prod = edm::makeRefToBaseProdFrom(jet_ref, iEvent);
0106     for (std::size_t i = 0; i < flav_names_.size(); i++) {
0107       output_tags.emplace_back(std::make_unique<JetTagCollection>(ref2prod));
0108     }
0109
0110     // init data storage
0111     data_.clear();
0112     for (const auto& len : input_sizes_) {
0113       data_.emplace_back(tag_infos->size() * len, 0.0);
0114     }
0115
0116     // convert inputs
0117     for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) {
0118       const auto& taginfo = (*tag_infos)[jet_n];
0119       make_inputs(jet_n, taginfo);
0120     }
0121
0122     // run prediction
0123     auto outputs = globalCache()->run(input_names_, data_, {}, output_names_, tag_infos->size())[0];
0124     assert(outputs.size() == flav_names_.size() * tag_infos->size());
0125
0126     // get the outputs
0127     unsigned i_output = 0;
0128     for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) {
0129       const auto& jet_ref = tag_infos->at(jet_n).jet();
0130       for (std::size_t flav_n = 0; flav_n < flav_names_.size(); flav_n++) {
0131         (*(output_tags[flav_n]))[jet_ref] = outputs[i_output];
0132         ++i_output;
0133       }
0134     }
0135   } else {
0136     // create empty output collection
0137     for (std::size_t i = 0; i < flav_names_.size(); i++) {
0138       output_tags.emplace_back(std::make_unique<JetTagCollection>());
0139     }
0140   }
0141
0142   // put into the event
0143   for (std::size_t flav_n = 0; flav_n < flav_names_.size(); ++flav_n) {
0144     iEvent.put(std::move(output_tags[flav_n]), flav_names_[flav_n]);
0145   }
0146 }
0147
0148 void DeepFlavourONNXJetTagsProducer::make_inputs(unsigned i_jet, const reco::DeepFlavourTagInfo& taginfo) {
0149   const auto& features = taginfo.features();
0150   float* ptr = nullptr;
0151   const float* start = nullptr;
0152   unsigned offset = 0;
0153
0154   // jet and other global features
0155   offset = i_jet * input_sizes_[kGlobal];
0156   ptr = &data_[kGlobal][offset];
0157   // jet variables
0158   const auto& jet_features = features.jet_features;
0159   start = ptr;
0160   *ptr = jet_features.pt;
0161   *(++ptr) = jet_features.eta;
0162   // number of elements in different collections
0163   *(++ptr) = features.c_pf_features.size();
0164   *(++ptr) = features.n_pf_features.size();
0165   *(++ptr) = features.sv_features.size();
0166   *(++ptr) = features.npv;
0167   // variables from ShallowTagInfo
0168   const auto& tag_info_features = features.tag_info_features;
0169   *(++ptr) = tag_info_features.trackSumJetEtRatio;
0170   *(++ptr) = tag_info_features.trackSumJetDeltaR;
0171   *(++ptr) = tag_info_features.vertexCategory;
0172   *(++ptr) = tag_info_features.trackSip2dValAboveCharm;
0173   *(++ptr) = tag_info_features.trackSip2dSigAboveCharm;
0174   *(++ptr) = tag_info_features.trackSip3dValAboveCharm;
0175   *(++ptr) = tag_info_features.trackSip3dSigAboveCharm;
0176   *(++ptr) = tag_info_features.jetNSelectedTracks;
0177   *(++ptr) = tag_info_features.jetNTracksEtaRel;
0178   assert(start + n_features_global_ - 1 == ptr);
0179
0180   // c_pf candidates
0181   auto max_c_pf_n = std::min(features.c_pf_features.size(), (std::size_t)25);
0182   offset = i_jet * input_sizes_[kChargedCandidates];
0183   for (std::size_t c_pf_n = 0; c_pf_n < max_c_pf_n; c_pf_n++) {
0184     const auto& c_pf_features = features.c_pf_features.at(c_pf_n);
0185     ptr = &data_[kChargedCandidates][offset + c_pf_n * n_features_cpf_];
0186     start = ptr;
0187     *ptr = c_pf_features.btagPf_trackEtaRel;
0188     *(++ptr) = c_pf_features.btagPf_trackPtRel;
0189     *(++ptr) = c_pf_features.btagPf_trackPPar;
0190     *(++ptr) = c_pf_features.btagPf_trackDeltaR;
0191     *(++ptr) = c_pf_features.btagPf_trackPParRatio;
0192     *(++ptr) = c_pf_features.btagPf_trackSip2dVal;
0193     *(++ptr) = c_pf_features.btagPf_trackSip2dSig;
0194     *(++ptr) = c_pf_features.btagPf_trackSip3dVal;
0195     *(++ptr) = c_pf_features.btagPf_trackSip3dSig;
0196     *(++ptr) = c_pf_features.btagPf_trackJetDistVal;
0197     *(++ptr) = c_pf_features.ptrel;
0198     *(++ptr) = c_pf_features.drminsv;
0199     *(++ptr) = c_pf_features.vtx_ass;
0200     *(++ptr) = c_pf_features.puppiw;
0201     *(++ptr) = c_pf_features.chi2;
0202     *(++ptr) = c_pf_features.quality;
0203     assert(start + n_features_cpf_ - 1 == ptr);
0204   }
0205
0206   // n_pf candidates
0207   auto max_n_pf_n = std::min(features.n_pf_features.size(), (std::size_t)25);
0208   offset = i_jet * input_sizes_[kNeutralCandidates];
0209   for (std::size_t n_pf_n = 0; n_pf_n < max_n_pf_n; n_pf_n++) {
0210     const auto& n_pf_features = features.n_pf_features.at(n_pf_n);
0211     ptr = &data_[kNeutralCandidates][offset + n_pf_n * n_features_npf_];
0212     start = ptr;
0213     *ptr = n_pf_features.ptrel;
0214     *(++ptr) = n_pf_features.deltaR;
0215     *(++ptr) = n_pf_features.isGamma;
0216     *(++ptr) = n_pf_features.hadFrac;
0217     *(++ptr) = n_pf_features.drminsv;
0218     *(++ptr) = n_pf_features.puppiw;
0219     assert(start + n_features_npf_ - 1 == ptr);
0220   }
0221
0222   // sv candidates
0223   auto max_sv_n = std::min(features.sv_features.size(), (std::size_t)4);
0224   offset = i_jet * input_sizes_[kVertices];
0225   for (std::size_t sv_n = 0; sv_n < max_sv_n; sv_n++) {
0226     const auto& sv_features = features.sv_features.at(sv_n);
0227     ptr = &data_[kVertices][offset + sv_n * n_features_sv_];
0228     start = ptr;
0229     *ptr = sv_features.pt;
0230     *(++ptr) = sv_features.deltaR;
0231     *(++ptr) = sv_features.mass;
0232     *(++ptr) = sv_features.ntracks;
0233     *(++ptr) = sv_features.chi2;
0234     *(++ptr) = sv_features.normchi2;
0235     *(++ptr) = sv_features.dxy;
0236     *(++ptr) = sv_features.dxysig;
0237     *(++ptr) = sv_features.d3d;
0238     *(++ptr) = sv_features.d3dsig;
0239     *(++ptr) = sv_features.costhetasvpv;
0240     *(++ptr) = sv_features.enratio;
0241     assert(start + n_features_sv_ - 1 == ptr);
0242   }
0243
0244   // last input: jet pt
0245   offset = i_jet * input_sizes_[kJetPt];
0246   data_[kJetPt][offset] = features.jet_features.pt;
0247 }
0248
0249 //define this as a plug-in
0250 DEFINE_FWK_MODULE(DeepFlavourONNXJetTagsProducer);