Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2022-06-30 23:17:13

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);