METPUSubtraction/plugins/DeepMETSonicProducer.cc

0001 #include "FWCore/Framework/interface/Event.h"
0002 #include "FWCore/Framework/interface/MakerMacros.h"
0003 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0004 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0005 #include "FWCore/Utilities/interface/InputTag.h"
0006 #include "DataFormats/PatCandidates/interface/MET.h"
0007 #include "DataFormats/PatCandidates/interface/PackedCandidate.h"
0008 #include "HeterogeneousCore/SonicTriton/interface/TritonEDProducer.h"
0009 #include "RecoMET/METPUSubtraction/interface/DeepMETHelp.h"
0010
0011 using namespace deepmet_helper;
0012
0013 class DeepMETSonicProducer : public TritonEDProducer<> {
0014 public:
0015   explicit DeepMETSonicProducer(const edm::ParameterSet&);
0016   void acquire(edm::Event const& iEvent, edm::EventSetup const& iSetup, Input& iInput) override;
0017   void produce(edm::Event& iEvent, edm::EventSetup const& iSetup, Output const& iOutput) override;
0018   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0019
0020 private:
0021   const edm::EDGetTokenT<std::vector<pat::PackedCandidate>> pf_token_;
0022   const float norm_;
0023   const bool ignore_leptons_;
0024   const unsigned int max_n_pf_;
0025   const float scale_;
0026   float px_leptons_;
0027   float py_leptons_;
0028 };
0029
0030 DeepMETSonicProducer::DeepMETSonicProducer(const edm::ParameterSet& cfg)
0031     : TritonEDProducer<>(cfg),
0032       pf_token_(consumes<std::vector<pat::PackedCandidate>>(cfg.getParameter<edm::InputTag>("pf_src"))),
0033       norm_(cfg.getParameter<double>("norm_factor")),
0034       ignore_leptons_(cfg.getParameter<bool>("ignore_leptons")),
0035       max_n_pf_(cfg.getParameter<unsigned int>("max_n_pf")),
0036       scale_(1.0 / norm_) {
0037   produces<pat::METCollection>();
0038 }
0039
0040 void DeepMETSonicProducer::acquire(edm::Event const& iEvent, edm::EventSetup const& iSetup, Input& iInput) {
0041   // one event per batch
0042   client_->setBatchSize(1);
0043   px_leptons_ = 0.;
0044   py_leptons_ = 0.;
0045
0046   auto const& pfs = iEvent.get(pf_token_);
0047
0048   auto& input = iInput.at("input");
0049   auto pfdata = input.allocate<float>();
0050   auto& vpfdata = (*pfdata)[0];
0051
0052   auto& input_cat0 = iInput.at("input_cat0");
0053   auto pfchg = input_cat0.allocate<float>();
0054   auto& vpfchg = (*pfchg)[0];
0055
0056   auto& input_cat1 = iInput.at("input_cat1");
0057   auto pfpdgId = input_cat1.allocate<float>();
0058   auto& vpfpdgId = (*pfpdgId)[0];
0059
0060   auto& input_cat2 = iInput.at("input_cat2");
0061   auto pffromPV = input_cat2.allocate<float>();
0062   auto& vpffromPV = (*pffromPV)[0];
0063
0064   size_t i_pf = 0;
0065   for (const auto& pf : pfs) {
0066     if (ignore_leptons_) {
0067       int pdg_id = std::abs(pf.pdgId());
0068       if (pdg_id == 11 || pdg_id == 13) {
0069         px_leptons_ += pf.px();
0070         py_leptons_ += pf.py();
0071         continue;
0072       }
0073     }
0074
0075     // PF keys [b'PF_dxy', b'PF_dz', b'PF_eta', b'PF_mass', b'PF_pt', b'PF_puppiWeight', b'PF_px', b'PF_py']
0076     vpfdata.push_back(rm_outlier(pf.dxy()));
0077     vpfdata.push_back(rm_outlier(pf.dz()));
0078     vpfdata.push_back(rm_outlier(pf.eta()));
0079     vpfdata.push_back(rm_outlier(pf.mass()));
0080     vpfdata.push_back(scale_and_rm_outlier(pf.pt(), scale_));
0081     vpfdata.push_back(rm_outlier(pf.puppiWeight()));
0082     vpfdata.push_back(scale_and_rm_outlier(pf.px(), scale_));
0083     vpfdata.push_back(scale_and_rm_outlier(pf.py(), scale_));
0084
0085     vpfchg.push_back(charge_embedding.at(pf.charge()));
0086
0087     vpfpdgId.push_back(pdg_id_embedding.at(pf.pdgId()));
0088
0089     vpffromPV.push_back(pf.fromPV());
0090
0091     ++i_pf;
0092     if (i_pf == max_n_pf_) {
0093       edm::LogWarning("acquire")
0094           << "<DeepMETSonicProducer::acquire>:" << std::endl
0095           << " The number of particles is equal to or exceeds the maximum considerable for DeepMET" << std::endl;
0096       break;
0097     }
0098   }
0099
0100   // fill the remaining with zeros
0101   // resize the vector to 4500 for zero-padding
0102   vpfdata.resize(8 * max_n_pf_);
0103   vpfchg.resize(max_n_pf_);
0104   vpfpdgId.resize(max_n_pf_);
0105   vpffromPV.resize(max_n_pf_);
0106
0107   input.toServer(pfdata);
0108   input_cat0.toServer(pfchg);
0109   input_cat1.toServer(pfpdgId);
0110   input_cat2.toServer(pffromPV);
0111 }
0112
0113 void DeepMETSonicProducer::produce(edm::Event& iEvent, edm::EventSetup const& iSetup, Output const& iOutput) {
0114   const auto& output1 = iOutput.begin()->second;
0115   const auto& outputs = output1.fromServer<float>();
0116
0117   // outputs are px and py
0118   float px = outputs[0][0] * norm_;
0119   float py = outputs[0][1] * norm_;
0120
0121   // subtract the lepton pt contribution
0122   px -= px_leptons_;
0123   py -= py_leptons_;
0124
0125   LogDebug("produce") << "<DeepMETSonicProducer::produce>:" << std::endl
0126                       << " MET from DeepMET Sonic Producer is MET_x " << px << " and MET_y " << py << std::endl;
0127
0128   auto pf_mets = std::make_unique<pat::METCollection>();
0129   const reco::Candidate::LorentzVector p4(px, py, 0., std::hypot(px, py));
0130   pf_mets->emplace_back(reco::MET(p4, {}));
0131   iEvent.put(std::move(pf_mets));
0132 }
0133
0134 void DeepMETSonicProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0135   edm::ParameterSetDescription desc;
0136   TritonClient::fillPSetDescription(desc);
0137   desc.add<edm::InputTag>("pf_src", edm::InputTag("packedPFCandidates"));
0138   desc.add<bool>("ignore_leptons", false);
0139   desc.add<double>("norm_factor", 50.);
0140   desc.add<unsigned int>("max_n_pf", 4500);
0141   descriptions.add("deepMETSonicProducer", desc);
0142 }
0143
0144 DEFINE_FWK_MODULE(DeepMETSonicProducer);