CMS 3D CMS Logo

DeepMETProducer.cc
Go to the documentation of this file.
5 
8 
11 
12 using namespace deepmet_helper;
13 
14 class DeepMETProducer : public edm::stream::EDProducer<edm::GlobalCache<tensorflow::SessionCache> > {
15 public:
17  void produce(edm::Event& event, const edm::EventSetup& setup) override;
18  static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
19 
20  // static methods for handling the global cache
21  static std::unique_ptr<tensorflow::SessionCache> initializeGlobalCache(const edm::ParameterSet&);
23 
24 private:
26  const float norm_;
27  const bool ignore_leptons_;
28  const unsigned int max_n_pf_;
29 
30  const tensorflow::Session* session_;
31 
32  tensorflow::Tensor input_;
33  tensorflow::Tensor input_cat0_;
34  tensorflow::Tensor input_cat1_;
35  tensorflow::Tensor input_cat2_;
36 };
37 
39  : pf_token_(consumes<std::vector<pat::PackedCandidate> >(cfg.getParameter<edm::InputTag>("pf_src"))),
40  norm_(cfg.getParameter<double>("norm_factor")),
41  ignore_leptons_(cfg.getParameter<bool>("ignore_leptons")),
42  max_n_pf_(cfg.getParameter<unsigned int>("max_n_pf")),
43  session_(cache->getSession()) {
44  produces<pat::METCollection>();
45 
46  const tensorflow::TensorShape shape({1, max_n_pf_, 8});
47  const tensorflow::TensorShape cat_shape({1, max_n_pf_, 1});
48 
49  input_ = tensorflow::Tensor(tensorflow::DT_FLOAT, shape);
50  input_cat0_ = tensorflow::Tensor(tensorflow::DT_FLOAT, cat_shape);
51  input_cat1_ = tensorflow::Tensor(tensorflow::DT_FLOAT, cat_shape);
52  input_cat2_ = tensorflow::Tensor(tensorflow::DT_FLOAT, cat_shape);
53 }
54 
56  auto const& pfs = event.get(pf_token_);
57 
58  const tensorflow::NamedTensorList input_list = {
59  {"input", input_}, {"input_cat0", input_cat0_}, {"input_cat1", input_cat1_}, {"input_cat2", input_cat2_}};
60 
61  // Set all inputs to zero
62  input_.flat<float>().setZero();
63  input_cat0_.flat<float>().setZero();
64  input_cat1_.flat<float>().setZero();
65  input_cat2_.flat<float>().setZero();
66 
67  size_t i_pf = 0;
68  float px_leptons = 0.;
69  float py_leptons = 0.;
70  const float scale = 1. / norm_;
71  for (const auto& pf : pfs) {
72  if (ignore_leptons_) {
73  int pdg_id = std::abs(pf.pdgId());
74  if (pdg_id == 11 || pdg_id == 13) {
75  px_leptons += pf.px();
76  py_leptons += pf.py();
77  continue;
78  }
79  }
80 
81  // fill the tensor
82  // 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']
83  float* ptr = &input_.tensor<float, 3>()(0, i_pf, 0);
84  *ptr = rm_outlier(pf.dxy());
85  *(++ptr) = rm_outlier(pf.dz());
86  *(++ptr) = rm_outlier(pf.eta());
87  *(++ptr) = rm_outlier(pf.mass());
88  *(++ptr) = scale_and_rm_outlier(pf.pt(), scale);
89  *(++ptr) = rm_outlier(pf.puppiWeight());
90  *(++ptr) = scale_and_rm_outlier(pf.px(), scale);
91  *(++ptr) = scale_and_rm_outlier(pf.py(), scale);
92  input_cat0_.tensor<float, 3>()(0, i_pf, 0) = charge_embedding.at(pf.charge());
93  input_cat1_.tensor<float, 3>()(0, i_pf, 0) = pdg_id_embedding.at(pf.pdgId());
94  input_cat2_.tensor<float, 3>()(0, i_pf, 0) = pf.fromPV();
95 
96  ++i_pf;
97  if (i_pf == max_n_pf_) {
98  break; // output a warning?
99  }
100  }
101 
102  std::vector<tensorflow::Tensor> outputs;
103  const std::vector<std::string> output_names = {"output/BiasAdd"};
104 
105  // run the inference and return met
107 
108  // The DNN directly estimates the missing px and py, not the recoil
109  float px = outputs[0].tensor<float, 2>()(0, 0) * norm_;
110  float py = outputs[0].tensor<float, 2>()(0, 1) * norm_;
111 
112  px -= px_leptons;
113  py -= py_leptons;
114 
115  LogDebug("produce") << "<DeepMETProducer::produce>:" << std::endl
116  << " MET from DeepMET Producer is MET_x " << px << " and MET_y " << py << std::endl;
117 
118  auto pf_mets = std::make_unique<pat::METCollection>();
120  pf_mets->emplace_back(reco::MET(p4, {}));
121  event.put(std::move(pf_mets));
122 }
123 
124 std::unique_ptr<tensorflow::SessionCache> DeepMETProducer::initializeGlobalCache(const edm::ParameterSet& params) {
125  // this method is supposed to create, initialize and return a SessionCache instance
126  std::string graphPath = edm::FileInPath(params.getParameter<std::string>("graph_path")).fullPath();
127  return std::make_unique<tensorflow::SessionCache>(graphPath);
128 }
129 
132  desc.add<edm::InputTag>("pf_src", edm::InputTag("packedPFCandidates"));
133  desc.add<bool>("ignore_leptons", false);
134  desc.add<double>("norm_factor", 50.);
135  desc.add<unsigned int>("max_n_pf", 4500);
136  desc.add<std::string>("graph_path", "RecoMET/METPUSubtraction/data/models/deepmet/deepmet_v1_2018/model.graphdef");
137  descriptions.add("deepMETProducer", desc);
138 }
139 
std::vector< NamedTensor > NamedTensorList
Definition: TensorFlow.h:31
MPlex< T, D1, D2, N > hypot(const MPlex< T, D1, D2, N > &a, const MPlex< T, D1, D2, N > &b)
Definition: Matriplex.h:616
const bool ignore_leptons_
void produce(edm::Event &event, const edm::EventSetup &setup) override
tensorflow::Tensor input_cat0_
tensorflow::Tensor input_cat1_
float rm_outlier(float val)
static void fillDescriptions(edm::ConfigurationDescriptions &descriptions)
tensorflow::Tensor input_
static std::unique_ptr< tensorflow::SessionCache > initializeGlobalCache(const edm::ParameterSet &)
static const std::unordered_map< int, int32_t > charge_embedding
Definition: DeepMETHelp.h:11
Definition: HeavyIon.h:7
Definition: MET.h:41
void run(Session *session, const NamedTensorList &inputs, const std::vector< std::string > &outputNames, std::vector< Tensor > *outputs, const thread::ThreadPoolOptions &threadPoolOptions)
Definition: TensorFlow.cc:271
Abs< T >::type abs(const T &t)
Definition: Abs.h:22
static void globalEndJob(tensorflow::SessionCache *)
#define DEFINE_FWK_MODULE(type)
Definition: MakerMacros.h:16
static const std::unordered_map< int, int32_t > pdg_id_embedding
Definition: DeepMETHelp.h:12
const tensorflow::Session * session_
static void fillDescriptions(edm::ConfigurationDescriptions &descriptions)
void add(std::string const &label, ParameterSetDescription const &psetDescription)
tensorflow::Tensor input_cat2_
math::XYZTLorentzVector LorentzVector
Lorentz vector.
Definition: Candidate.h:36
DeepMETProducer(const edm::ParameterSet &, const tensorflow::SessionCache *)
HLT enums.
def cache(function)
Definition: utilities.py:3
float scale_and_rm_outlier(float val, float scale)
Definition: DeepMETHelper.cc:5
const unsigned int max_n_pf_
const std::string & fullPath() const
Definition: FileInPath.cc:144
const edm::EDGetTokenT< std::vector< pat::PackedCandidate > > pf_token_
def move(src, dest)
Definition: eostools.py:511
Definition: event.py:1
const float norm_
#define LogDebug(id)