Skip to content

Commit 56fe406

Browse files
cmsbuildcmsbuild
authored
Merge pull request #49429 from sbein/from-CMSSW_16_0_X_2025-11-18-2300
Refined pT-based sorting of FastSim jets in NANO, Type-1 MET correction
2 parents ba95571 + 6c9384a commit 56fe406

File tree

3 files changed

+262
-33
lines changed

3 files changed

+262
-33
lines changed

‎PhysicsTools/NanoAOD/plugins/JetSorter.cc‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,57 @@
1+
#include "FWCore/Framework/interface/global/EDProducer.h"
2+
#include "FWCore/Framework/interface/Event.h"
3+
#include "FWCore/Framework/interface/MakerMacros.h"
4+
#include "FWCore/ParameterSet/interface/ParameterSet.h"
5+
6+
#include "DataFormats/Common/interface/View.h"
7+
#include "DataFormats/PatCandidates/interface/Jet.h"
8+
9+
#include <algorithm>
10+
#include <string>
11+
#include <memory>
12+
13+
class JetSorter : public edm::global::EDProducer<> {
14+
public:
15+
explicit JetSorter(const edm::ParameterSet &iConfig)
16+
: srcToken_(consumes<edm::View<pat::Jet>>(iConfig.getParameter<edm::InputTag>("src"))),
17+
userFloatSorter_(iConfig.getParameter<std::string>("userFloatSorter")),
18+
descending_(iConfig.getParameter<bool>("descending")) {
19+
produces<std::vector<pat::Jet>>();
20+
}
21+
22+
~JetSorter() override = default;
23+
24+
void produce(edm::StreamID, edm::Event &iEvent, const edm::EventSetup &) const override {
25+
edm::Handle<edm::View<pat::Jet>> hSrc;
26+
iEvent.getByToken(srcToken_, hSrc);
27+
28+
auto out = std::make_unique<std::vector<pat::Jet>>();
29+
out->reserve(hSrc->size());
30+
31+
for (auto const &jIn : *hSrc) {
32+
out->push_back(jIn);
33+
}
34+
35+
auto key = [this](pat::Jet const &j) -> float {
36+
if (j.hasUserFloat(userFloatSorter_)) {
37+
return j.userFloat(userFloatSorter_);
38+
}
39+
return static_cast<float>(j.pt());
40+
};
41+
42+
std::stable_sort(out->begin(), out->end(), [this, &key](pat::Jet const &a, pat::Jet const &b) {
43+
float pa = key(a);
44+
float pb = key(b);
45+
return descending_ ? (pa > pb) : (pa < pb);
46+
});
47+
48+
iEvent.put(std::move(out));
49+
}
50+
51+
private:
52+
edm::EDGetTokenT<edm::View<pat::Jet>> srcToken_;
53+
std::string userFloatSorter_;
54+
bool descending_;
55+
};
56+
57+
DEFINE_FWK_MODULE(JetSorter);

‎PhysicsTools/NanoAOD/plugins/ProcessRefinedJets.cc‎

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,136 @@
1+
#include "FWCore/Framework/interface/global/EDProducer.h"
2+
#include "FWCore/Framework/interface/Event.h"
3+
#include "FWCore/Framework/interface/MakerMacros.h"
4+
#include "FWCore/ParameterSet/interface/ParameterSet.h"
5+
6+
#include "DataFormats/Common/interface/View.h"
7+
#include "DataFormats/PatCandidates/interface/Jet.h"
8+
#include "DataFormats/PatCandidates/interface/MET.h"
9+
10+
#include <cmath>
11+
#include <memory>
12+
#include <string>
13+
14+
//This producer embeds the final refined FastSim jet pT values
15+
//so that they can be used as input to the pT-based sorting routine, as well
16+
//as to implement the Type-1 MET correction based on the refined jets.
17+
class ProcessRefinedJets : public edm::global::EDProducer<> {
18+
public:
19+
explicit ProcessRefinedJets(const edm::ParameterSet &iConfig)
20+
: jetsToken_(consumes<edm::View<pat::Jet>>(iConfig.getParameter<edm::InputTag>("jets"))),
21+
refinedPtName_(iConfig.getParameter<std::string>("refinedPtName")),
22+
maskBtagName_(iConfig.getParameter<std::string>("maskBtagName")),
23+
ptFinalName_(iConfig.getParameter<std::string>("ptFinalName")),
24+
ptUnrefinedName_(iConfig.getParameter<std::string>("ptUnrefinedName")) {
25+
// Type-1 MET correction is optional
26+
if (iConfig.existsAs<edm::InputTag>("met")) {
27+
metToken_ = consumes<edm::View<pat::MET>>(iConfig.getParameter<edm::InputTag>("met"));
28+
doMET_ = true;
29+
} else {
30+
doMET_ = false;
31+
}
32+
33+
produces<std::vector<pat::Jet>>();
34+
35+
if (doMET_) {
36+
// refined MET made available in python
37+
produces<std::vector<pat::MET>>("Refined");
38+
}
39+
}
40+
41+
~ProcessRefinedJets() override = default;
42+
43+
void produce(edm::StreamID, edm::Event &iEvent, const edm::EventSetup &) const override {
44+
edm::Handle<edm::View<pat::Jet>> hJets;
45+
iEvent.getByToken(jetsToken_, hJets);
46+
47+
auto outJets = std::make_unique<std::vector<pat::Jet>>();
48+
outJets->reserve(hJets->size());
49+
50+
// Sum over (pt_orig - pt_final) for MET correction
51+
double sumDeltaPx = 0.;
52+
double sumDeltaPy = 0.;
53+
54+
for (auto const &jIn : *hJets) {
55+
pat::Jet j(jIn);
56+
57+
const double pt_orig = j.pt();
58+
const double phi = j.phi();
59+
60+
// mask: BvsAll > 0
61+
const float bvsAll = j.bDiscriminator(maskBtagName_);
62+
const bool refine = (bvsAll > 0.f);
63+
64+
const double pt_ref = j.hasUserFloat(refinedPtName_) ? static_cast<double>(j.userFloat(refinedPtName_)) : pt_orig;
65+
66+
const double pt_final = refine ? pt_ref : pt_orig;
67+
68+
// store unrefined pt if requested
69+
if (!ptUnrefinedName_.empty()) {
70+
j.addUserFloat(ptUnrefinedName_, static_cast<float>(pt_orig));
71+
}
72+
73+
// store final pt used for Nano, needed for the sorting
74+
j.addUserFloat(ptFinalName_, static_cast<float>(pt_final));
75+
76+
// MET delta: add unrefined jets, subtract final jets
77+
// This is equivalent to adding (pt_orig - pt_final) in the jet direction.
78+
const double dpt = pt_orig - pt_final;
79+
sumDeltaPx += dpt * std::cos(phi);
80+
sumDeltaPy += dpt * std::sin(phi);
81+
82+
outJets->push_back(std::move(j));
83+
}
84+
85+
iEvent.put(std::move(outJets));
86+
87+
// Optionally correct MET
88+
if (doMET_) {
89+
edm::Handle<edm::View<pat::MET>> hMET;
90+
iEvent.getByToken(metToken_, hMET);
91+
92+
auto outMET = std::make_unique<std::vector<pat::MET>>();
93+
outMET->reserve(hMET->size());
94+
95+
for (auto const &mIn : *hMET) {
96+
pat::MET m(mIn);
97+
98+
const double pxOrig = m.px();
99+
const double pyOrig = m.py();
100+
const double ptOrig = m.pt();
101+
const double phiOrig = m.phi();
102+
103+
const double pxFinal = pxOrig + sumDeltaPx;
104+
const double pyFinal = pyOrig + sumDeltaPy;
105+
const double ptFinal = std::sqrt(pxFinal * pxFinal + pyFinal * pyFinal);
106+
const double phiFinal = std::atan2(pyFinal, pxFinal);
107+
108+
m.addUserFloat("pt_unrefined", static_cast<float>(ptOrig));
109+
m.addUserFloat("phi_unrefined", static_cast<float>(phiOrig));
110+
111+
// stash final values, needed for sorting.
112+
m.addUserFloat("pt_final", static_cast<float>(ptFinal));
113+
m.addUserFloat("phi_final", static_cast<float>(phiFinal));
114+
115+
// push back into pat::MET with reco::Candidate::LorentzVector ---
116+
reco::Candidate::LorentzVector p4(pxFinal, pyFinal, 0.0, ptFinal);
117+
m.setP4(p4);
118+
119+
outMET->push_back(std::move(m));
120+
}
121+
iEvent.put(std::move(outMET), "Refined");
122+
}
123+
}
124+
125+
private:
126+
edm::EDGetTokenT<edm::View<pat::Jet>> jetsToken_;
127+
edm::EDGetTokenT<edm::View<pat::MET>> metToken_;
128+
bool doMET_ = false;
129+
130+
std::string refinedPtName_;
131+
std::string maskBtagName_;
132+
std::string ptFinalName_;
133+
std::string ptUnrefinedName_;
134+
};
135+
136+
DEFINE_FWK_MODULE(ProcessRefinedJets);

‎PhysicsTools/NanoAOD/python/jetsAK4_Puppi_cff.py‎

Lines changed: 69 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -220,22 +220,11 @@ def nanoAOD_addDeepInfoAK4(process,addParticleNet,addRobustParTAK4=False,addUnif
220220
jetPuppiTablesTask = cms.Task(jetPuppiTable)
221221

222222
from Configuration.Eras.Modifier_fastSim_cff import fastSim
223-
from PhysicsTools.NanoAOD.common_cff import ExtVar
224-
225-
import FWCore.ParameterSet.Config as cms
226-
from Configuration.Eras.Modifier_fastSim_cff import fastSim
227-
from PhysicsTools.NanoAOD.common_cff import Var
228-
229223
from PhysicsTools.NanoAOD.common_cff import Var, ExtVar
230-
from Configuration.Eras.Modifier_fastSim_cff import fastSim
231-
232-
from Configuration.Eras.Modifier_fastSim_cff import fastSim
233-
from PhysicsTools.NanoAOD.common_cff import ExtVar, Var
234-
import FWCore.ParameterSet.Config as cms
235224

236225
def nanoAOD_refineFastSim_puppiJet(process):
237226

238-
#save originals and clear so we can republish
227+
# 0. Save originals and clear so we can republish refined versions
239228
fastSim.toModify(process.jetPuppiTable.variables,
240229
pt_unrefined = process.jetPuppiTable.variables.pt.clone(),
241230
btagDeepFlavB_unrefined = process.jetPuppiTable.variables.btagDeepFlavB.clone(),
@@ -252,13 +241,13 @@ def nanoAOD_refineFastSim_puppiJet(process):
252241
btagUParTAK4B=None, btagUParTAK4CvB=None, btagUParTAK4CvL=None, btagUParTAK4QvG=None,
253242
)
254243

255-
# run refinement model and return features
244+
# 1. Run refinement model and return features
256245
process.puppiJetRefineNN = cms.EDProducer(
257246
"JetBaseMVAValueMapProducer",
258247
backend = cms.string("ONNX"),
259248
batch_eval = cms.bool(True),
260249
disableONNXGraphOpt = cms.bool(True),
261-
src = cms.InputTag("linkedObjects","jets"), # <— matches table
250+
src = cms.InputTag("linkedObjects","jets"), # matches table
262251
weightFile = cms.FileInPath("PhysicsTools/NanoAOD/data/fastSimPuppiJetRefineNN_31July2025.onnx"),
263252
name = cms.string("puppiJetRefineNN"),
264253
variables = cms.VPSet(
@@ -273,25 +262,24 @@ def nanoAOD_refineFastSim_puppiJet(process):
273262
cms.PSet(name=cms.string("Jet_btagUParTAK4B"), expr=cms.string("bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:BvsAll')")),
274263
cms.PSet(name=cms.string("Jet_btagUParTAK4CvB"), expr=cms.string("?(bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:CvsB')>0)?bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:CvsB'):-1")),
275264
cms.PSet(name=cms.string("Jet_btagUParTAK4CvL"), expr=cms.string("?(bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:CvsL')>0)?bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:CvsL'):-1")),
276-
cms.PSet(name=cms.string("Jet_btagUParTAK4QvG"), expr=cms.string("?(bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:QvsG')>0)?bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:QvsG'):-1")),
277-
),
265+
cms.PSet(name=cms.string("Jet_btagUParTAK4QvG"), expr=cms.string("?(bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:QvsG')>0)?bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:QvsG'):-1")),),
278266
inputTensorName = cms.string("input"),
279267
outputTensorName = cms.string("output"),
280268
outputNames = cms.vstring(
281269
"ptrefined",
282270
"btagDeepFlavBrefined","btagDeepFlavCvBrefined","btagDeepFlavCvLrefined","btagDeepFlavQGrefined",
283-
"btagUParTAK4Brefined","btagUParTAK4CvBrefined","btagUParTAK4CvLrefined","btagUParTAK4QvGrefined",
284-
),
271+
"btagUParTAK4Brefined","btagUParTAK4CvBrefined","btagUParTAK4CvLrefined","btagUParTAK4QvGrefined",),
285272
outputFormulas = cms.vstring("at(0)","at(1)","at(2)","at(3)","at(4)","at(5)","at(6)","at(7)","at(8)"),
286273
)
287274
fastSim.toModify(process.jetPuppiTablesTask, process.jetPuppiTablesTask.add(process.puppiJetRefineNN))
288275

289-
#copy the ONNX ValueMaps onto jets as userFloats
276+
# Ensure src is what we expect (redundant but explicit)
290277
process.puppiJetRefineNN.src = cms.InputTag("linkedObjects","jets")
291-
278+
279+
# 2. Copy the ONNX ValueMaps onto jets as userFloats
292280
process.finalJetsPuppiWithRefined = cms.EDProducer(
293281
"PATJetUserDataEmbedder",
294-
src = cms.InputTag("linkedObjects","jets"), # <— same!
282+
src = cms.InputTag("linkedObjects","jets"),
295283
userFloats = cms.PSet(
296284
ptrefined = cms.InputTag("puppiJetRefineNN","ptrefined"),
297285
btagDeepFlavBrefined = cms.InputTag("puppiJetRefineNN","btagDeepFlavBrefined"),
@@ -304,26 +292,74 @@ def nanoAOD_refineFastSim_puppiJet(process):
304292
btagUParTAK4QvGrefined = cms.InputTag("puppiJetRefineNN","btagUParTAK4QvGrefined"),
305293
)
306294
)
307-
process.jetPuppiTable.src = cms.InputTag("finalJetsPuppiWithRefined")
308-
309295
fastSim.toModify(process.jetPuppiTablesTask, process.jetPuppiTablesTask.add(process.finalJetsPuppiWithRefined))
310296

297+
# Intermediate src: jets with refined userFloats
311298
process.jetPuppiTable.src = cms.InputTag("finalJetsPuppiWithRefined")
312299

313-
#mask jets we shouldn't be refining - Note: the -1 safety in the definitions of the original flavor taggers cannot be implemented in the following, because nested ternaries either crash the parser or confuse the logic and mess up the default values. The max(x,-1) is used as an alternative.
300+
# 3. Apply mask for all refined quantities in Nano (pt + taggers)
301+
# Note: we keep all the masking logic here in python.
314302
_mask = "bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:BvsAll')>0"
303+
315304
fastSim.toModify(process.jetPuppiTable.variables,
316-
pt = Var(f"?{_mask}?userFloat('ptrefined'):pt()", float, precision=10),
317-
btagDeepFlavB = Var(f"?{_mask}?userFloat('btagDeepFlavBrefined'):(bDiscriminator('pfDeepFlavourJetTags:probb')+bDiscriminator('pfDeepFlavourJetTags:probbb')+bDiscriminator('pfDeepFlavourJetTags:problepb'))", float, precision=10),
318-
btagDeepFlavCvB = Var(f"?{_mask}?userFloat('btagDeepFlavCvBrefined'):max(bDiscriminator('pfDeepFlavourJetTags:probc')/(bDiscriminator('pfDeepFlavourJetTags:probc')+bDiscriminator('pfDeepFlavourJetTags:probb')+bDiscriminator('pfDeepFlavourJetTags:probbb')+bDiscriminator('pfDeepFlavourJetTags:problepb')),-1)", float, precision=10),#max(x,-1) safety because double-ternary not allowed
319-
btagDeepFlavCvL = Var(f"?{_mask}?userFloat('btagDeepFlavCvLrefined'):max(bDiscriminator('pfDeepFlavourJetTags:probc')/(bDiscriminator('pfDeepFlavourJetTags:probc')+bDiscriminator('pfDeepFlavourJetTags:probuds')+bDiscriminator('pfDeepFlavourJetTags:probg')),-1)", float, precision=10),#max(x,-1) safety because double-ternary not allowed
320-
btagDeepFlavQG = Var(f"?{_mask}?userFloat('btagDeepFlavQGrefined'):max(bDiscriminator('pfDeepFlavourJetTags:probg')/(bDiscriminator('pfDeepFlavourJetTags:probg')+bDiscriminator('pfDeepFlavourJetTags:probuds')),-1)", float, precision=10),#max(x,-1) safety because double-ternary not allowed
321-
btagUParTAK4B = Var(f"?{_mask}?userFloat('btagUParTAK4Brefined'):max(bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:BvsAll'),-1)", float, precision=12),
322-
btagUParTAK4CvB = Var(f"?{_mask}?userFloat('btagUParTAK4CvBrefined'):max(bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:CvsB'),-1)", float, precision=12),
323-
btagUParTAK4CvL = Var(f"?{_mask}?userFloat('btagUParTAK4CvLrefined'):max(bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:CvsL'),-1)", float, precision=12),
324-
btagUParTAK4QvG = Var(f"?{_mask}?userFloat('btagUParTAK4QvGrefined'):max(bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:QvsG'),-1)", float, precision=12),
305+
pt = Var("?" + _mask + "?userFloat('ptrefined'):pt()", float, precision=10),
306+
btagDeepFlavB = Var("?" + _mask + "?userFloat('btagDeepFlavBrefined'):(bDiscriminator('pfDeepFlavourJetTags:probb')+bDiscriminator('pfDeepFlavourJetTags:probbb')+bDiscriminator('pfDeepFlavourJetTags:problepb'))", float, precision=10),
307+
btagDeepFlavCvB = Var("?" + _mask + "?userFloat('btagDeepFlavCvBrefined'):max(bDiscriminator('pfDeepFlavourJetTags:probc')/(bDiscriminator('pfDeepFlavourJetTags:probc')+bDiscriminator('pfDeepFlavourJetTags:probb')+bDiscriminator('pfDeepFlavourJetTags:probbb')+bDiscriminator('pfDeepFlavourJetTags:problepb')),-1)", float, precision=10),
308+
btagDeepFlavCvL = Var("?" + _mask + "?userFloat('btagDeepFlavCvLrefined'):max(bDiscriminator('pfDeepFlavourJetTags:probc')/(bDiscriminator('pfDeepFlavourJetTags:probc')+bDiscriminator('pfDeepFlavourJetTags:probuds')+bDiscriminator('pfDeepFlavourJetTags:probg')),-1)", float, precision=10),
309+
btagDeepFlavQG = Var("?" + _mask + "?userFloat('btagDeepFlavQGrefined'):max(bDiscriminator('pfDeepFlavourJetTags:probg')/(bDiscriminator('pfDeepFlavourJetTags:probg')+bDiscriminator('pfDeepFlavourJetTags:probuds')),-1)", float, precision=10),
310+
btagUParTAK4B = Var("?" + _mask + "?userFloat('btagUParTAK4Brefined'):max(bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:BvsAll'),-1)", float, precision=12),
311+
btagUParTAK4CvB = Var("?" + _mask + "?userFloat('btagUParTAK4CvBrefined'):max(bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:CvsB'),-1)", float, precision=12),
312+
btagUParTAK4CvL = Var("?" + _mask + "?userFloat('btagUParTAK4CvLrefined'):max(bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:CvsL'),-1)", float, precision=12),
313+
btagUParTAK4QvG = Var("?" + _mask + "?userFloat('btagUParTAK4QvGrefined'):max(bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:QvsG'),-1)", float, precision=12),
314+
)
315+
316+
# 4. Build pt_final as a userFloat on jets (mask applied), also do type 1 met correction
317+
# Note for future refinement: The mask definition above is also coded into the producer
318+
process.processRefinedJets = cms.EDProducer(
319+
"ProcessRefinedJets", # or "FastSimPuppiRefinedJetProducer" if that is your C++ class name
320+
jets = cms.InputTag("finalJetsPuppiWithRefined"),
321+
refinedPtName = cms.string("ptrefined"),
322+
maskBtagName = cms.string("pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:BvsAll"),
323+
ptFinalName = cms.string("pt_final"),
324+
ptUnrefinedName = cms.string("pt_unrefined"),
325+
met = cms.InputTag("slimmedMETsPuppi"), # MET collection where type-1 refinement corrections
325326
)
327+
fastSim.toModify(
328+
process.jetPuppiTablesTask,
329+
process.jetPuppiTablesTask.add(process.processRefinedJets)
330+
)
331+
332+
#point PuppiMET Nano table to the *refined* MET collection
333+
fastSim.toModify(
334+
process.puppiMetTable,
335+
src = cms.InputTag("processRefinedJets", "Refined"),
336+
)
337+
338+
# And add _unrefined branches from the MET userFloats filled in C++
339+
fastSim.toModify(
340+
process.puppiMetTable.variables,
341+
pt_unrefined = Var("userFloat('pt_unrefined')", float, precision=10),
342+
phi_unrefined = Var("userFloat('phi_unrefined')", float, precision=10),
343+
)
344+
345+
# 5. Sort jets by pt_final only
346+
process.finalJetsPuppiSorted = cms.EDProducer(
347+
"JetSorter",
348+
src = cms.InputTag("processRefinedJets"),
349+
userFloatSorter = cms.string("pt_final"),
350+
descending = cms.bool(True),
351+
)
352+
fastSim.toModify(
353+
process.jetPuppiTablesTask,
354+
process.jetPuppiTablesTask.add(process.finalJetsPuppiSorted)
355+
)
356+
fastSim.toModify(
357+
process.jetPuppiTable,
358+
src = cms.InputTag("finalJetsPuppiSorted")
359+
)
360+
326361
return process
327362

363+
328364
# bind for the customise step
329365
nanoAOD_refineFastSim_puppiJet = nanoAOD_refineFastSim_puppiJet

0 commit comments

Comments
 (0)