ONNXRuntime/plugins/ParticleNetSonicJetTagsProducer.cc

0001 #include "FWCore/Framework/interface/Frameworkfwd.h"
0002
0003 #include "FWCore/Framework/interface/Event.h"
0004 #include "FWCore/Framework/interface/MakerMacros.h"
0005
0006 #include "FWCore/Framework/interface/makeRefToBaseProdFrom.h"
0007
0008 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0009 #include "FWCore/Utilities/interface/StreamID.h"
0010
0011 #include "DataFormats/BTauReco/interface/JetTag.h"
0012
0013 #include "DataFormats/BTauReco/interface/DeepBoostedJetTagInfo.h"
0014
0015 #include "HeterogeneousCore/SonicTriton/interface/TritonEDProducer.h"
0016
0017 #include "HeterogeneousCore/SonicTriton/interface/TritonData.h"
0018
0019 #include "RecoBTag/FeatureTools/interface/deep_helpers.h"
0020
0021 #include <iostream>
0022 #include <fstream>
0023 #include <algorithm>
0024 #include <numeric>
0025 #include <nlohmann/json.hpp>
0026
0027 using namespace btagbtvdeep;
0028
0029 class ParticleNetSonicJetTagsProducer : public TritonEDProducer<> {
0030 public:
0031   explicit ParticleNetSonicJetTagsProducer(const edm::ParameterSet &);
0032   ~ParticleNetSonicJetTagsProducer() override;
0033
0034   void acquire(edm::Event const &iEvent, edm::EventSetup const &iSetup, Input &iInput) override;
0035   void produce(edm::Event &iEvent, edm::EventSetup const &iSetup, Output const &iOutput) override;
0036   static void fillDescriptions(edm::ConfigurationDescriptions &);
0037
0038 private:
0039   typedef std::vector<reco::DeepBoostedJetTagInfo> TagInfoCollection;
0040   typedef reco::JetTagCollection JetTagCollection;
0041
0042   void make_inputs(const reco::DeepBoostedJetTagInfo &taginfo);
0043
0044   const edm::EDGetTokenT<TagInfoCollection> src_;
0045   std::vector<std::string> flav_names_;             // names of the output scores
0046   std::vector<std::string> input_names_;            // names of each input group - the ordering is important!
0047   std::vector<std::vector<int64_t>> input_shapes_;  // shapes of each input group (-1 for dynamic axis)
0048   std::vector<unsigned> input_sizes_;               // total length of each input vector
0049   std::unordered_map<std::string, PreprocessParams> prep_info_map_;  // preprocessing info for each input group
0050   bool debug_ = false;
0051   bool skippedInference_ = false;
0052   constexpr static unsigned numParticleGroups_ = 3;
0053 };
0054
0055 ParticleNetSonicJetTagsProducer::ParticleNetSonicJetTagsProducer(const edm::ParameterSet &iConfig)
0056     : TritonEDProducer<>(iConfig),
0057       src_(consumes<TagInfoCollection>(iConfig.getParameter<edm::InputTag>("src"))),
0058       flav_names_(iConfig.getParameter<std::vector<std::string>>("flav_names")),
0059       debug_(iConfig.getUntrackedParameter<bool>("debugMode", false)) {
0060   ParticleNetConstructor(iConfig, false, input_names_, prep_info_map_, input_shapes_, input_sizes_, nullptr);
0061
0062   if (debug_) {
0063     LogDebug("ParticleNetSonicJetTagsProducer") << "<ParticleNetSonicJetTagsProducer::produce>:" << std::endl;
0064     for (unsigned i = 0; i < input_names_.size(); ++i) {
0065       const auto &group_name = input_names_.at(i);
0066       if (!input_shapes_.empty()) {
0067         LogDebug("ParticleNetSonicJetTagsProducer") << group_name << "\nshapes: ";
0068         for (const auto &x : input_shapes_.at(i)) {
0069           LogDebug("ParticleNetSonicJetTagsProducer") << x << ", ";
0070         }
0071       }
0072       LogDebug("ParticleNetSonicJetTagsProducer") << "\nvariables: ";
0073       for (const auto &x : prep_info_map_.at(group_name).var_names) {
0074         LogDebug("ParticleNetSonicJetTagsProducer") << x << ", ";
0075       }
0076       LogDebug("ParticleNetSonicJetTagsProducer") << "\n";
0077     }
0078     LogDebug("ParticleNetSonicJetTagsProducer") << "flav_names: ";
0079     for (const auto &flav_name : flav_names_) {
0080       LogDebug("ParticleNetSonicJetTagsProducer") << flav_name << ", ";
0081     }
0082     LogDebug("ParticleNetSonicJetTagsProducer") << "\n";
0083   }
0084
0085   // get output names from flav_names
0086   for (const auto &flav_name : flav_names_) {
0087     produces<JetTagCollection>(flav_name);
0088   }
0089 }
0090
0091 ParticleNetSonicJetTagsProducer::~ParticleNetSonicJetTagsProducer() {}
0092
0093 void ParticleNetSonicJetTagsProducer::fillDescriptions(edm::ConfigurationDescriptions &descriptions) {
0094   // pfDeepBoostedJetTags
0095   edm::ParameterSetDescription desc;
0096   TritonClient::fillPSetDescription(desc);
0097   desc.add<edm::InputTag>("src", edm::InputTag("pfDeepBoostedJetTagInfos"));
0098   desc.add<std::string>("preprocess_json", "");
0099   // `preprocessParams` is deprecated -- use the preprocessing json file instead
0100   edm::ParameterSetDescription preprocessParams;
0101   preprocessParams.setAllowAnything();
0102   preprocessParams.setComment("`preprocessParams` is deprecated, please use `preprocess_json` instead.");
0103   desc.addOptional<edm::ParameterSetDescription>("preprocessParams", preprocessParams);
0104   desc.add<std::vector<std::string>>("flav_names",
0105                                      std::vector<std::string>{
0106                                          "probTbcq",
0107                                          "probTbqq",
0108                                          "probTbc",
0109                                          "probTbq",
0110                                          "probWcq",
0111                                          "probWqq",
0112                                          "probZbb",
0113                                          "probZcc",
0114                                          "probZqq",
0115                                          "probHbb",
0116                                          "probHcc",
0117                                          "probHqqqq",
0118                                          "probQCDbb",
0119                                          "probQCDcc",
0120                                          "probQCDb",
0121                                          "probQCDc",
0122                                          "probQCDothers",
0123                                      });
0124   desc.addOptionalUntracked<bool>("debugMode", false);
0125
0126   descriptions.addWithDefaultLabel(desc);
0127 }
0128
0129 void ParticleNetSonicJetTagsProducer::acquire(edm::Event const &iEvent, edm::EventSetup const &iSetup, Input &iInput) {
0130   edm::Handle<TagInfoCollection> tag_infos;
0131   iEvent.getByToken(src_, tag_infos);
0132   client_->setBatchSize(tag_infos->size());
0133   skippedInference_ = false;
0134   if (!tag_infos->empty()) {
0135     unsigned int maxParticles = 0;
0136     unsigned int maxVertices = 0;
0137     for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) {
0138       maxParticles = std::max(maxParticles,
0139                               static_cast<unsigned int>(((*tag_infos)[jet_n]).features().get("pfcand_etarel").size()));
0140       maxVertices =
0141           std::max(maxVertices, static_cast<unsigned int>(((*tag_infos)[jet_n]).features().get("sv_etarel").size()));
0142     }
0143     if (maxParticles == 0 && maxVertices == 0) {
0144       client_->setBatchSize(0);
0145       skippedInference_ = true;
0146       return;
0147     }
0148     unsigned int minPartFromJSON = prep_info_map_.at(input_names_[0]).min_length;
0149     unsigned int maxPartFromJSON = prep_info_map_.at(input_names_[0]).max_length;
0150     unsigned int minVertFromJSON = prep_info_map_.at(input_names_[3]).min_length;
0151     unsigned int maxVertFromJSON = prep_info_map_.at(input_names_[3]).max_length;
0152     maxParticles = std::clamp(maxParticles, minPartFromJSON, maxPartFromJSON);
0153     maxVertices = std::clamp(maxVertices, minVertFromJSON, maxVertFromJSON);
0154
0155     for (unsigned igroup = 0; igroup < input_names_.size(); ++igroup) {
0156       const auto &group_name = input_names_[igroup];
0157       auto &input = iInput.at(group_name);
0158       unsigned target;
0159       if (igroup < numParticleGroups_) {
0160         input.setShape(1, maxParticles);
0161         target = maxParticles;
0162       } else {
0163         input.setShape(1, maxVertices);
0164         target = maxVertices;
0165       }
0166       auto tdata = input.allocate<float>(true);
0167       for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) {
0168         const auto &taginfo = (*tag_infos)[jet_n];
0169         auto &vdata = (*tdata)[jet_n];
0170         const auto &prep_params = prep_info_map_.at(group_name);
0171         unsigned curr_pos = 0;
0172         // transform/pad
0173         for (unsigned i = 0; i < prep_params.var_names.size(); ++i) {
0174           const auto &varname = prep_params.var_names[i];
0175           const auto &raw_value = taginfo.features().get(varname);
0176           const auto &info = prep_params.info(varname);
0177           int insize = center_norm_pad_halfRagged(raw_value,
0178                                                   info.center,
0179                                                   info.norm_factor,
0180                                                   target,
0181                                                   vdata,
0182                                                   curr_pos,
0183                                                   info.pad,
0184                                                   info.replace_inf_value,
0185                                                   info.lower_bound,
0186                                                   info.upper_bound);
0187           curr_pos += insize;
0188           if (i == 0 && (!input_shapes_.empty())) {
0189             input_shapes_[igroup][2] = insize;
0190           }
0191
0192           if (debug_) {
0193             LogDebug("acquire") << "<ParticleNetSonicJetTagsProducer::produce>:" << std::endl
0194                                 << " -- var=" << varname << ", center=" << info.center << ", scale=" << info.norm_factor
0195                                 << ", replace=" << info.replace_inf_value << ", pad=" << info.pad << std::endl;
0196             for (unsigned i = curr_pos - insize; i < curr_pos; i++) {
0197               LogDebug("acquire") << vdata[i] << ",";
0198             }
0199             LogDebug("acquire") << std::endl;
0200           }
0201         }
0202       }
0203       input.toServer(tdata);
0204     }
0205   }
0206 }
0207
0208 void ParticleNetSonicJetTagsProducer::produce(edm::Event &iEvent,
0209                                               const edm::EventSetup &iSetup,
0210                                               Output const &iOutput) {
0211   edm::Handle<TagInfoCollection> tag_infos;
0212   iEvent.getByToken(src_, tag_infos);
0213
0214   // initialize output collection
0215   std::vector<std::unique_ptr<JetTagCollection>> output_tags;
0216   if (!tag_infos->empty()) {
0217     auto jet_ref = tag_infos->begin()->jet();
0218     auto ref2prod = edm::makeRefToBaseProdFrom(jet_ref, iEvent);
0219     for (std::size_t i = 0; i < flav_names_.size(); i++) {
0220       output_tags.emplace_back(std::make_unique<JetTagCollection>(ref2prod));
0221     }
0222   } else {
0223     for (std::size_t i = 0; i < flav_names_.size(); i++) {
0224       output_tags.emplace_back(std::make_unique<JetTagCollection>());
0225     }
0226   }
0227
0228   if (!tag_infos->empty()) {
0229     if (!skippedInference_) {
0230       const auto &output1 = iOutput.begin()->second;
0231       const auto &outputs_from_server = output1.fromServer<float>();
0232
0233       for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) {
0234         const auto &taginfo = (*tag_infos)[jet_n];
0235         const auto &jet_ref = tag_infos->at(jet_n).jet();
0236
0237         if (!taginfo.features().empty()) {
0238           for (std::size_t flav_n = 0; flav_n < flav_names_.size(); flav_n++) {
0239             (*(output_tags[flav_n]))[jet_ref] = outputs_from_server[jet_n][flav_n];
0240           }
0241         } else {
0242           for (std::size_t flav_n = 0; flav_n < flav_names_.size(); flav_n++) {
0243             (*(output_tags[flav_n]))[jet_ref] = 0.;
0244           }
0245         }
0246       }
0247     } else {
0248       for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) {
0249         const auto &jet_ref = tag_infos->at(jet_n).jet();
0250         for (std::size_t flav_n = 0; flav_n < flav_names_.size(); flav_n++) {
0251           (*(output_tags[flav_n]))[jet_ref] = 0.;
0252         }
0253       }
0254     }
0255   }
0256
0257   if (debug_) {
0258     LogDebug("produce") << "<ParticleNetSonicJetTagsProducer::produce>:" << std::endl
0259                         << "=== " << iEvent.id().run() << ":" << iEvent.id().luminosityBlock() << ":"
0260                         << iEvent.id().event() << " ===" << std::endl;
0261     for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) {
0262       const auto &jet_ref = tag_infos->at(jet_n).jet();
0263       LogDebug("produce") << " - Jet #" << jet_n << ", pt=" << jet_ref->pt() << ", eta=" << jet_ref->eta()
0264                           << ", phi=" << jet_ref->phi() << std::endl;
0265       for (std::size_t flav_n = 0; flav_n < flav_names_.size(); ++flav_n) {
0266         LogDebug("produce") << "    " << flav_names_.at(flav_n) << " = " << (*(output_tags.at(flav_n)))[jet_ref]
0267                             << std::endl;
0268       }
0269     }
0270   }
0271
0272   // put into the event
0273   for (std::size_t flav_n = 0; flav_n < flav_names_.size(); ++flav_n) {
0274     iEvent.put(std::move(output_tags[flav_n]), flav_names_[flav_n]);
0275   }
0276 }
0277
0278 //define this as a plug-in
0279 DEFINE_FWK_MODULE(ParticleNetSonicJetTagsProducer);