Adding DeepMET SONIC Producer to CMSSW

Configuration/ProcessModifiers/python/allSonicTriton_cff.py

@@ -3,6 +3,7 @@
 from Configuration.ProcessModifiers.enableSonicTriton_cff import enableSonicTriton
 from Configuration.ProcessModifiers.particleNetSonicTriton_cff import particleNetSonicTriton
 from Configuration.ProcessModifiers.particleNetPTSonicTriton_cff import particleNetPTSonicTriton
+from Configuration.ProcessModifiers.deepMETSonicTriton_cff import deepMETSonicTriton
 
 # collect all SonicTriton-related process modifiers here
-allSonicTriton = cms.ModifierChain(enableSonicTriton,particleNetSonicTriton)
+allSonicTriton = cms.ModifierChain(enableSonicTriton,deepMETSonicTriton,particleNetSonicTriton)

Configuration/ProcessModifiers/python/deepMETSonicTriton_cff.py

@@ -0,0 +1,3 @@
+import FWCore.ParameterSet.Config as cms
+
+deepMETSonicTriton = cms.Modifier()

PhysicsTools/NanoAOD/python/nano_cff.py

@@ -156,10 +156,10 @@ def nanoAOD_addBoostedTauIds(process, idsToRun=[]):
 def nanoAOD_recalibrateMETs(process,isData):
     # add DeepMETs
     nanoAOD_DeepMET_switch = cms.PSet(
-        ResponseTune_Graph = cms.untracked.string('RecoMET/METPUSubtraction/data/deepmet/deepmet_resp_v1_2018.pb')
+        ResponseTune_Graph = cms.untracked.string('RecoMET/METPUSubtraction/data/models/deepmet/deepmet_resp_v1_2018/model.graphdef')
     )
     for modifier in run2_miniAOD_80XLegacy, run2_nanoAOD_94X2016:
-        modifier.toModify(nanoAOD_DeepMET_switch, ResponseTune_Graph=cms.untracked.string("RecoMET/METPUSubtraction/data/deepmet/deepmet_resp_v1_2016.pb"))
+        modifier.toModify(nanoAOD_DeepMET_switch, ResponseTune_Graph="RecoMET/METPUSubtraction/data/models/deepmet/deepmet_resp_v1_2016/model.graphdef")
 
     print("add DeepMET Producers")
     process.load('RecoMET.METPUSubtraction.deepMETProducer_cfi')

PhysicsTools/PatAlgos/python/slimming/miniAOD_tools.py

@@ -494,29 +494,10 @@ def _add_slimmedMETsPuppi(process):
     (~pp_on_AA).toModify(process, _add_slimmedMETsPuppi)
 
     def _add_deepMET(process):
-        process.load('RecoMET.METPUSubtraction.deepMETProducer_cfi')
+        from RecoMET.METPUSubtraction.deepMETProducer_cff import deepMETsResolutionTune, deepMETsResponseTune
 
-        addToProcessAndTask('deepMETsResolutionTune', process.deepMETProducer.clone(), process, task)
-        addToProcessAndTask('deepMETsResponseTune', process.deepMETProducer.clone(), process, task)
-        process.deepMETsResponseTune.graph_path = 'RecoMET/METPUSubtraction/data/deepmet/deepmet_resp_v1_2018.pb'
-
-        from Configuration.Eras.Modifier_phase2_common_cff import phase2_common
-        phase2_common.toModify(
-            process.deepMETsResolutionTune,
-            max_n_pf=12500,
-            graph_path="RecoMET/METPUSubtraction/data/deepmet/deepmet_v1_phase2.pb"
-        )
-        phase2_common.toModify(
-            process.deepMETsResponseTune,
-            max_n_pf=12500,
-            graph_path="RecoMET/METPUSubtraction/data/deepmet/deepmet_resp_v1_phase2.pb"
-        )
-
-        from Configuration.Eras.Modifier_run2_jme_2016_cff import run2_jme_2016
-        run2_jme_2016.toModify(
-            process.deepMETsResponseTune,
-            graph_path="RecoMET/METPUSubtraction/data/deepmet/deepmet_resp_v1_2016.pb"
-        )
+        addToProcessAndTask('deepMETsResolutionTune', deepMETsResolutionTune, process, task)
+        addToProcessAndTask('deepMETsResponseTune', deepMETsResponseTune, process, task)
     (~pp_on_AA).toModify(process, _add_deepMET)
 
     # add DetIdAssociatorRecords to EventSetup (for isolatedTracks)

RecoMET/METPUSubtraction/interface/DeepMETHelp.h

@@ -0,0 +1,15 @@
+#ifndef RecoMET_METPUSubtraction_DeepMETHelp_h
+#define RecoMET_METPUSubtraction_DeepMETHelp_h
+
+#include <unordered_map>
+#include <cstdint>
+
+namespace deepmet_helper {
+  float scale_and_rm_outlier(float val, float scale);
+
+  static const std::unordered_map<int, int32_t> charge_embedding{{-1, 0}, {0, 1}, {1, 2}};
+  static const std::unordered_map<int, int32_t> pdg_id_embedding{
+      {-211, 0}, {-13, 1}, {-11, 2}, {0, 3}, {1, 4}, {2, 5}, {11, 6}, {13, 7}, {22, 8}, {130, 9}, {211, 10}};
+}  // namespace deepmet_helper
+
+#endif

RecoMET/METPUSubtraction/plugins/BuildFile.xml

@@ -16,6 +16,7 @@
   <use name="RecoMET/METPUSubtraction"/>
   <use name="RecoJets/JetProducers"/>
   <use name="PhysicsTools/TensorFlow"/>
+  <use name="HeterogeneousCore/SonicTriton"/>
   <use name="root"/>
   <use name="roottmva"/>
 </library>

RecoMET/METPUSubtraction/plugins/DeepMETProducer.cc

@@ -7,6 +7,9 @@
 #include "DataFormats/PatCandidates/interface/PackedCandidate.h"
 
 #include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
+#include "RecoMET/METPUSubtraction/interface/DeepMETHelp.h"
+
+using namespace deepmet_helper;
 
 struct DeepMETCache {
   std::atomic<tensorflow::GraphDef*> graph_def;
@@ -34,21 +37,8 @@ class DeepMETProducer : public edm::stream::EDProducer<edm::GlobalCache<DeepMETC
   tensorflow::Tensor input_cat0_;
   tensorflow::Tensor input_cat1_;
   tensorflow::Tensor input_cat2_;
-
-  inline static const std::unordered_map<int, int32_t> charge_embedding_{{-1, 0}, {0, 1}, {1, 2}};
-  inline static const std::unordered_map<int, int32_t> pdg_id_embedding_{
-      {-211, 0}, {-13, 1}, {-11, 2}, {0, 3}, {1, 4}, {2, 5}, {11, 6}, {13, 7}, {22, 8}, {130, 9}, {211, 10}};
 };
 
-namespace {
-  float scale_and_rm_outlier(float val, float scale) {
-    float ret_val = val * scale;
-    if (ret_val > 1e6 || ret_val < -1e6)
-      return 0.;
-    return ret_val;
-  }
-}  // namespace
-
 DeepMETProducer::DeepMETProducer(const edm::ParameterSet& cfg, const DeepMETCache* cache)
     : pf_token_(consumes<std::vector<pat::PackedCandidate> >(cfg.getParameter<edm::InputTag>("pf_src"))),
       norm_(cfg.getParameter<double>("norm_factor")),
@@ -103,8 +93,8 @@ void DeepMETProducer::produce(edm::Event& event, const edm::EventSetup& setup) {
     *(++ptr) = pf.puppiWeight();
     *(++ptr) = scale_and_rm_outlier(pf.px(), scale);
     *(++ptr) = scale_and_rm_outlier(pf.py(), scale);
-    input_cat0_.tensor<float, 3>()(0, i_pf, 0) = charge_embedding_.at(pf.charge());
-    input_cat1_.tensor<float, 3>()(0, i_pf, 0) = pdg_id_embedding_.at(pf.pdgId());
+    input_cat0_.tensor<float, 3>()(0, i_pf, 0) = charge_embedding.at(pf.charge());
+    input_cat1_.tensor<float, 3>()(0, i_pf, 0) = pdg_id_embedding.at(pf.pdgId());
     input_cat2_.tensor<float, 3>()(0, i_pf, 0) = pf.fromPV();
 
     ++i_pf;
@@ -126,6 +116,9 @@ void DeepMETProducer::produce(edm::Event& event, const edm::EventSetup& setup) {
   px -= px_leptons;
   py -= py_leptons;
 
+  LogDebug("produce") << "<DeepMETProducer::produce>:" << std::endl
+                      << " MET from DeepMET Producer is MET_x " << px << " and MET_y " << py << std::endl;
+
   auto pf_mets = std::make_unique<pat::METCollection>();
   const reco::Candidate::LorentzVector p4(px, py, 0., std::hypot(px, py));
   pf_mets->emplace_back(reco::MET(p4, {}));
@@ -154,7 +147,7 @@ void DeepMETProducer::fillDescriptions(edm::ConfigurationDescriptions& descripti
   desc.add<bool>("ignore_leptons", false);
   desc.add<double>("norm_factor", 50.);
   desc.add<unsigned int>("max_n_pf", 4500);
-  desc.add<std::string>("graph_path", "RecoMET/METPUSubtraction/data/deepmet/deepmet_v1_2018.pb");
+  desc.add<std::string>("graph_path", "RecoMET/METPUSubtraction/data/models/deepmet/deepmet_v1_2018/model.graphdef");
   descriptions.add("deepMETProducer", desc);
 }
 

RecoMET/METPUSubtraction/plugins/DeepMETSonicProducer.cc

@@ -0,0 +1,144 @@
+#include "FWCore/Framework/interface/Event.h"
+#include "FWCore/Framework/interface/MakerMacros.h"
+#include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
+#include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
+#include "FWCore/Utilities/interface/InputTag.h"
+#include "DataFormats/PatCandidates/interface/MET.h"
+#include "DataFormats/PatCandidates/interface/PackedCandidate.h"
+#include "HeterogeneousCore/SonicTriton/interface/TritonEDProducer.h"
+#include "RecoMET/METPUSubtraction/interface/DeepMETHelp.h"
+
+using namespace deepmet_helper;
+
+class DeepMETSonicProducer : public TritonEDProducer<> {
+public:
+  explicit DeepMETSonicProducer(const edm::ParameterSet&);
+  void acquire(edm::Event const& iEvent, edm::EventSetup const& iSetup, Input& iInput) override;
+  void produce(edm::Event& iEvent, edm::EventSetup const& iSetup, Output const& iOutput) override;
+  static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
+
+private:
+  const edm::EDGetTokenT<std::vector<pat::PackedCandidate>> pf_token_;
+  const float norm_;
+  const bool ignore_leptons_;
+  const unsigned int max_n_pf_;
+  const float scale_;
+  float px_leptons_;
+  float py_leptons_;
+};
+
+DeepMETSonicProducer::DeepMETSonicProducer(const edm::ParameterSet& cfg)
+    : TritonEDProducer<>(cfg),
+      pf_token_(consumes<std::vector<pat::PackedCandidate>>(cfg.getParameter<edm::InputTag>("pf_src"))),
+      norm_(cfg.getParameter<double>("norm_factor")),
+      ignore_leptons_(cfg.getParameter<bool>("ignore_leptons")),
+      max_n_pf_(cfg.getParameter<unsigned int>("max_n_pf")),
+      scale_(1.0 / norm_) {
+  produces<pat::METCollection>();
+}
+
+void DeepMETSonicProducer::acquire(edm::Event const& iEvent, edm::EventSetup const& iSetup, Input& iInput) {
+  // one event per batch
+  client_->setBatchSize(1);
+  px_leptons_ = 0.;
+  py_leptons_ = 0.;
+
+  auto const& pfs = iEvent.get(pf_token_);
+
+  auto& input = iInput.at("input");
+  auto pfdata = input.allocate<float>();
+  auto& vpfdata = (*pfdata)[0];
+
+  auto& input_cat0 = iInput.at("input_cat0");
+  auto pfchg = input_cat0.allocate<float>();
+  auto& vpfchg = (*pfchg)[0];
+
+  auto& input_cat1 = iInput.at("input_cat1");
+  auto pfpdgId = input_cat1.allocate<float>();
+  auto& vpfpdgId = (*pfpdgId)[0];
+
+  auto& input_cat2 = iInput.at("input_cat2");
+  auto pffromPV = input_cat2.allocate<float>();
+  auto& vpffromPV = (*pffromPV)[0];
+
+  size_t i_pf = 0;
+  for (const auto& pf : pfs) {
+    if (ignore_leptons_) {
+      int pdg_id = std::abs(pf.pdgId());
+      if (pdg_id == 11 || pdg_id == 13) {
+        px_leptons_ += pf.px();
+        py_leptons_ += pf.py();
+        continue;
+      }
+    }
+
+    // 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']
+    vpfdata.push_back(pf.dxy());
+    vpfdata.push_back(pf.dz());
+    vpfdata.push_back(pf.eta());
+    vpfdata.push_back(pf.mass());
+    vpfdata.push_back(scale_and_rm_outlier(pf.pt(), scale_));
+    vpfdata.push_back(pf.puppiWeight());
+    vpfdata.push_back(scale_and_rm_outlier(pf.px(), scale_));
+    vpfdata.push_back(scale_and_rm_outlier(pf.py(), scale_));
+
+    vpfchg.push_back(charge_embedding.at(pf.charge()));
+
+    vpfpdgId.push_back(pdg_id_embedding.at(pf.pdgId()));
+
+    vpffromPV.push_back(pf.fromPV());
+
+    ++i_pf;
+    if (i_pf == max_n_pf_) {
+      edm::LogWarning("acquire")
+          << "<DeepMETSonicProducer::acquire>:" << std::endl
+          << " The number of particles is equal to or exceeds the maximum considerable for DeepMET" << std::endl;
+      break;
+    }
+  }
+
+  // fill the remaining with zeros
+  // resize the vector to 4500 for zero-padding
+  vpfdata.resize(8 * max_n_pf_);
+  vpfchg.resize(max_n_pf_);
+  vpfpdgId.resize(max_n_pf_);
+  vpffromPV.resize(max_n_pf_);
+
+  input.toServer(pfdata);
+  input_cat0.toServer(pfchg);
+  input_cat1.toServer(pfpdgId);
+  input_cat2.toServer(pffromPV);
+}
+
+void DeepMETSonicProducer::produce(edm::Event& iEvent, edm::EventSetup const& iSetup, Output const& iOutput) {
+  const auto& output1 = iOutput.begin()->second;
+  const auto& outputs = output1.fromServer<float>();
+
+  // outputs are px and py
+  float px = outputs[0][0] * norm_;
+  float py = outputs[0][1] * norm_;
+
+  // subtract the lepton pt contribution
+  px -= px_leptons_;
+  py -= py_leptons_;
+
+  LogDebug("produce") << "<DeepMETSonicProducer::produce>:" << std::endl
+                      << " MET from DeepMET Sonic Producer is MET_x " << px << " and MET_y " << py << std::endl;
+
+  auto pf_mets = std::make_unique<pat::METCollection>();
+  const reco::Candidate::LorentzVector p4(px, py, 0., std::hypot(px, py));
+  pf_mets->emplace_back(reco::MET(p4, {}));
+  iEvent.put(std::move(pf_mets));
+}
+
+void DeepMETSonicProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
+  edm::ParameterSetDescription desc;
+  TritonClient::fillPSetDescription(desc);
+  desc.add<edm::InputTag>("pf_src", edm::InputTag("packedPFCandidates"));
+  desc.add<bool>("ignore_leptons", false);
+  desc.add<double>("norm_factor", 50.);
+  desc.add<unsigned int>("max_n_pf", 4500);
+  descriptions.add("deepMETSonicProducer", desc);
+}
+
+DEFINE_FWK_MODULE(DeepMETSonicProducer);

RecoMET/METPUSubtraction/python/deepMETProducer_cff.py

@@ -0,0 +1,55 @@
+import FWCore.ParameterSet.Config as cms
+from FWCore.ParameterSet.pfnInPath import pfnInPath
+import os
+
+from RecoMET.METPUSubtraction.deepMETProducer_cfi import deepMETProducer as _deepMETProducer
+
+deepMETsResolutionTune = _deepMETProducer.clone()
+deepMETsResponseTune = _deepMETProducer.clone(
+    graph_path = 'RecoMET/METPUSubtraction/data/models/deepmet/deepmet_resp_v1_2018/model.graphdef',
+)
+
+from Configuration.Eras.Modifier_phase2_common_cff import phase2_common
+phase2_common.toModify(
+    deepMETsResolutionTune,
+    max_n_pf=12500,
+    graph_path="RecoMET/METPUSubtraction/data/models/deepmet_phase2/deepmet_v1_phase2/model.graphdef"
+)
+phase2_common.toModify(
+    deepMETsResponseTune,
+    max_n_pf=12500,
+    graph_path="RecoMET/METPUSubtraction/data/models/deepmet_phase2/deepmet_resp_v1_phase2/model.graphdef"
+)
+
+from Configuration.Eras.Modifier_run2_jme_2016_cff import run2_jme_2016
+run2_jme_2016.toModify(
+    deepMETsResponseTune,
+    graph_path="RecoMET/METPUSubtraction/data/models/deepmet/deepmet_resp_v1_2016/model.graphdef"
+)
+
+from RecoMET.METPUSubtraction.deepMETSonicProducer_cff import deepMETSonicProducer as _deepMETSonicProducer
+from Configuration.ProcessModifiers.deepMETSonicTriton_cff import deepMETSonicTriton
+
+def split_model_path(path):
+    Client = dict(
+        modelName = path.split('/')[-3],
+        modelConfigPath = '/'.join(path.split('/')[:-2])+'/config.pbtxt',
+        # version "1" is the resolutionTune
+        # version "2" is the responeTune
+        modelVersion = os.path.realpath(pfnInPath(path).split(':')[-1]).split('/')[-2],
+    )
+    return Client
+
+# propagate possible parameter updates
+_deepMETsResolutionTuneSonic = _deepMETSonicProducer.clone(
+    max_n_pf = deepMETsResolutionTune.max_n_pf,
+    Client = split_model_path(deepMETsResolutionTune.graph_path.value()),
+)
+deepMETSonicTriton.toReplaceWith(deepMETsResolutionTune, _deepMETsResolutionTuneSonic)
+
+_deepMETsResponseTuneSonic = _deepMETSonicProducer.clone(
+    max_n_pf = deepMETsResponseTune.max_n_pf,
+    Client = split_model_path(deepMETsResponseTune.graph_path.value()),
+)
+deepMETSonicTriton.toReplaceWith(deepMETsResponseTune, _deepMETsResponseTuneSonic)
+

RecoMET/METPUSubtraction/python/deepMETSonicProducer_cff.py

@@ -0,0 +1,16 @@
+import FWCore.ParameterSet.Config as cms
+
+from RecoMET.METPUSubtraction.deepMETSonicProducer_cfi import deepMETSonicProducer as _deepMETSonicProducer
+
+deepMETSonicProducer = _deepMETSonicProducer.clone(
+    Client = dict(
+        timeout = 300,
+        mode = "Async",
+        modelName = "deepmet",
+        modelConfigPath = "RecoMET/METPUSubtraction/data/models/deepmet/config.pbtxt",
+        # version "1" is the resolutionTune
+        # version "2" is the responeTune
+        modelVersion = "1",
+    ),
+    pf_src = "packedPFCandidates",
+)

RecoMET/METPUSubtraction/src/DeepMETHelper.cc

@@ -0,0 +1,11 @@
+#include "RecoMET/METPUSubtraction/interface/DeepMETHelp.h"
+
+namespace deepmet_helper {
+  float scale_and_rm_outlier(float val, float scale) {
+    float ret_val = val * scale;
+    if (ret_val > 1e6 || ret_val < -1e6)
+      return 0.;
+    return ret_val;
+  }
+
+}  // namespace deepmet_helper