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_;
0046 std::vector<std::string> input_names_;
0047 std::vector<std::vector<int64_t>> input_shapes_;
0048 std::vector<unsigned> input_sizes_;
0049 std::unordered_map<std::string, PreprocessParams> prep_info_map_;
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
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
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
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
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
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
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
0279 DEFINE_FWK_MODULE(ParticleNetSonicJetTagsProducer);