Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2022-07-16 00:28:52

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