METPUSubtraction/plugins/DeepMETProducer.cc

0001 #include "FWCore/Framework/interface/Frameworkfwd.h"
0002 #include "FWCore/Framework/interface/stream/EDProducer.h"
0003 #include "FWCore/Framework/interface/Event.h"
0004 #include "FWCore/Framework/interface/MakerMacros.h"
0005
0006 #include "DataFormats/PatCandidates/interface/MET.h"
0007 #include "DataFormats/PatCandidates/interface/PackedCandidate.h"
0008
0009 #include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
0010 #include "RecoMET/METPUSubtraction/interface/DeepMETHelp.h"
0011
0012 using namespace deepmet_helper;
0013
0014 class DeepMETProducer : public edm::stream::EDProducer<edm::GlobalCache<tensorflow::SessionCache> > {
0015 public:
0016   explicit DeepMETProducer(const edm::ParameterSet&, const tensorflow::SessionCache*);
0017   void produce(edm::Event& event, const edm::EventSetup& setup) override;
0018   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0019
0020   // static methods for handling the global cache
0021   static std::unique_ptr<tensorflow::SessionCache> initializeGlobalCache(const edm::ParameterSet&);
0022   static void globalEndJob(tensorflow::SessionCache*){};
0023
0024 private:
0025   const edm::EDGetTokenT<std::vector<pat::PackedCandidate> > pf_token_;
0026   const float norm_;
0027   const bool ignore_leptons_;
0028   const unsigned int max_n_pf_;
0029
0030   const tensorflow::Session* session_;
0031
0032   tensorflow::Tensor input_;
0033   tensorflow::Tensor input_cat0_;
0034   tensorflow::Tensor input_cat1_;
0035   tensorflow::Tensor input_cat2_;
0036 };
0037
0038 DeepMETProducer::DeepMETProducer(const edm::ParameterSet& cfg, const tensorflow::SessionCache* cache)
0039     : pf_token_(consumes<std::vector<pat::PackedCandidate> >(cfg.getParameter<edm::InputTag>("pf_src"))),
0040       norm_(cfg.getParameter<double>("norm_factor")),
0041       ignore_leptons_(cfg.getParameter<bool>("ignore_leptons")),
0042       max_n_pf_(cfg.getParameter<unsigned int>("max_n_pf")),
0043       session_(cache->getSession()) {
0044   produces<pat::METCollection>();
0045
0046   const tensorflow::TensorShape shape({1, max_n_pf_, 8});
0047   const tensorflow::TensorShape cat_shape({1, max_n_pf_, 1});
0048
0049   input_ = tensorflow::Tensor(tensorflow::DT_FLOAT, shape);
0050   input_cat0_ = tensorflow::Tensor(tensorflow::DT_FLOAT, cat_shape);
0051   input_cat1_ = tensorflow::Tensor(tensorflow::DT_FLOAT, cat_shape);
0052   input_cat2_ = tensorflow::Tensor(tensorflow::DT_FLOAT, cat_shape);
0053 }
0054
0055 void DeepMETProducer::produce(edm::Event& event, const edm::EventSetup& setup) {
0056   auto const& pfs = event.get(pf_token_);
0057
0058   const tensorflow::NamedTensorList input_list = {
0059       {"input", input_}, {"input_cat0", input_cat0_}, {"input_cat1", input_cat1_}, {"input_cat2", input_cat2_}};
0060
0061   // Set all inputs to zero
0062   input_.flat<float>().setZero();
0063   input_cat0_.flat<float>().setZero();
0064   input_cat1_.flat<float>().setZero();
0065   input_cat2_.flat<float>().setZero();
0066
0067   size_t i_pf = 0;
0068   float px_leptons = 0.;
0069   float py_leptons = 0.;
0070   const float scale = 1. / norm_;
0071   for (const auto& pf : pfs) {
0072     if (ignore_leptons_) {
0073       int pdg_id = std::abs(pf.pdgId());
0074       if (pdg_id == 11 || pdg_id == 13) {
0075         px_leptons += pf.px();
0076         py_leptons += pf.py();
0077         continue;
0078       }
0079     }
0080
0081     // fill the tensor
0082     // 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']
0083     float* ptr = &input_.tensor<float, 3>()(0, i_pf, 0);
0084     *ptr = rm_outlier(pf.dxy());
0085     *(++ptr) = rm_outlier(pf.dz());
0086     *(++ptr) = rm_outlier(pf.eta());
0087     *(++ptr) = rm_outlier(pf.mass());
0088     *(++ptr) = scale_and_rm_outlier(pf.pt(), scale);
0089     *(++ptr) = rm_outlier(pf.puppiWeight());
0090     *(++ptr) = scale_and_rm_outlier(pf.px(), scale);
0091     *(++ptr) = scale_and_rm_outlier(pf.py(), scale);
0092     input_cat0_.tensor<float, 3>()(0, i_pf, 0) = charge_embedding.at(pf.charge());
0093     input_cat1_.tensor<float, 3>()(0, i_pf, 0) = pdg_id_embedding.at(pf.pdgId());
0094     input_cat2_.tensor<float, 3>()(0, i_pf, 0) = pf.fromPV();
0095
0096     ++i_pf;
0097     if (i_pf == max_n_pf_) {
0098       break;  // output a warning?
0099     }
0100   }
0101
0102   std::vector<tensorflow::Tensor> outputs;
0103   const std::vector<std::string> output_names = {"output/BiasAdd"};
0104
0105   // run the inference and return met
0106   tensorflow::run(session_, input_list, output_names, &outputs);
0107
0108   // The DNN directly estimates the missing px and py, not the recoil
0109   float px = outputs[0].tensor<float, 2>()(0, 0) * norm_;
0110   float py = outputs[0].tensor<float, 2>()(0, 1) * norm_;
0111
0112   px -= px_leptons;
0113   py -= py_leptons;
0114
0115   LogDebug("produce") << "<DeepMETProducer::produce>:" << std::endl
0116                       << " MET from DeepMET Producer is MET_x " << px << " and MET_y " << py << std::endl;
0117
0118   auto pf_mets = std::make_unique<pat::METCollection>();
0119   const reco::Candidate::LorentzVector p4(px, py, 0., std::hypot(px, py));
0120   pf_mets->emplace_back(reco::MET(p4, {}));
0121   event.put(std::move(pf_mets));
0122 }
0123
0124 std::unique_ptr<tensorflow::SessionCache> DeepMETProducer::initializeGlobalCache(const edm::ParameterSet& params) {
0125   // this method is supposed to create, initialize and return a SessionCache instance
0126   std::string graphPath = edm::FileInPath(params.getParameter<std::string>("graph_path")).fullPath();
0127   return std::make_unique<tensorflow::SessionCache>(graphPath);
0128 }
0129
0130 void DeepMETProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0131   edm::ParameterSetDescription desc;
0132   desc.add<edm::InputTag>("pf_src", edm::InputTag("packedPFCandidates"));
0133   desc.add<bool>("ignore_leptons", false);
0134   desc.add<double>("norm_factor", 50.);
0135   desc.add<unsigned int>("max_n_pf", 4500);
0136   desc.add<std::string>("graph_path", "RecoMET/METPUSubtraction/data/models/deepmet/deepmet_v1_2018/model.graphdef");
0137   descriptions.add("deepMETProducer", desc);
0138 }
0139
0140 DEFINE_FWK_MODULE(DeepMETProducer);