Skip to content

Commit

Permalink
Merge pull request #37963 from wpmccormack/CMSSW_12_4_0_pre2_deepMET_…
Browse files Browse the repository at this point in the history
…SONIC

Adding DeepMET SONIC Producer to CMSSW
  • Loading branch information
cmsbuild committed Jul 15, 2022
2 parents c1e18ac + 8c3bac9 commit a70261c
Show file tree
Hide file tree
Showing 12 changed files with 320 additions and 41 deletions.
3 changes: 2 additions & 1 deletion Configuration/ProcessModifiers/python/allSonicTriton_cff.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
import FWCore.ParameterSet.Config as cms

deepMETSonicTriton = cms.Modifier()
4 changes: 2 additions & 2 deletions PhysicsTools/NanoAOD/python/nano_cff.py
Original file line number Diff line number Diff line change
Expand Up @@ -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')
Expand Down
25 changes: 3 additions & 22 deletions PhysicsTools/PatAlgos/python/slimming/miniAOD_tools.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
15 changes: 15 additions & 0 deletions RecoMET/METPUSubtraction/interface/DeepMETHelp.h
Original file line number Diff line number Diff line change
@@ -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
1 change: 1 addition & 0 deletions RecoMET/METPUSubtraction/plugins/BuildFile.xml
Original file line number Diff line number Diff line change
Expand Up @@ -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>
25 changes: 9 additions & 16 deletions RecoMET/METPUSubtraction/plugins/DeepMETProducer.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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")),
Expand Down Expand Up @@ -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;
Expand All @@ -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, {}));
Expand Down Expand Up @@ -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);
}

Expand Down
144 changes: 144 additions & 0 deletions RecoMET/METPUSubtraction/plugins/DeepMETSonicProducer.cc
Original file line number Diff line number Diff line change
@@ -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);
55 changes: 55 additions & 0 deletions RecoMET/METPUSubtraction/python/deepMETProducer_cff.py
Original file line number Diff line number Diff line change
@@ -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)

16 changes: 16 additions & 0 deletions RecoMET/METPUSubtraction/python/deepMETSonicProducer_cff.py
Original file line number Diff line number Diff line change
@@ -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",
)
11 changes: 11 additions & 0 deletions RecoMET/METPUSubtraction/src/DeepMETHelper.cc
Original file line number Diff line number Diff line change
@@ -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
Loading

0 comments on commit a70261c

Please sign in to comment.