Skip to content

Commit 882bbc4

Browse files
committed
implemented refined pT-based sorting of jets in NANO, Type-1 MET correction of the refined jets
1 parent 318ca4d commit 882bbc4

File tree

3 files changed

+265
-35
lines changed

3 files changed

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

PhysicsTools/NanoAOD/python/jetsAK4_Puppi_cff.py

Lines changed: 67 additions & 35 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"),
@@ -301,29 +289,73 @@ def nanoAOD_refineFastSim_puppiJet(process):
301289
btagUParTAK4Brefined = cms.InputTag("puppiJetRefineNN","btagUParTAK4Brefined"),
302290
btagUParTAK4CvBrefined = cms.InputTag("puppiJetRefineNN","btagUParTAK4CvBrefined"),
303291
btagUParTAK4CvLrefined = cms.InputTag("puppiJetRefineNN","btagUParTAK4CvLrefined"),
304-
btagUParTAK4QvGrefined = cms.InputTag("puppiJetRefineNN","btagUParTAK4QvGrefined"),
305-
)
292+
btagUParTAK4QvGrefined = cms.InputTag("puppiJetRefineNN","btagUParTAK4QvGrefined"),)
306293
)
307-
process.jetPuppiTable.src = cms.InputTag("finalJetsPuppiWithRefined")
308-
309294
fastSim.toModify(process.jetPuppiTablesTask, process.jetPuppiTablesTask.add(process.finalJetsPuppiWithRefined))
310295

296+
# Intermediate src: jets with refined userFloats
311297
process.jetPuppiTable.src = cms.InputTag("finalJetsPuppiWithRefined")
312298

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.
299+
# 3. Apply mask for all refined quantities in Nano (pt + taggers)
300+
# Note: we keep all the masking logic here in python.
314301
_mask = "bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:BvsAll')>0"
302+
315303
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),
304+
pt = Var("?" + _mask + "?userFloat('ptrefined'):pt()", float, precision=10),
305+
btagDeepFlavB = Var("?" + _mask + "?userFloat('btagDeepFlavBrefined'):(bDiscriminator('pfDeepFlavourJetTags:probb')+bDiscriminator('pfDeepFlavourJetTags:probbb')+bDiscriminator('pfDeepFlavourJetTags:problepb'))", float, precision=10),
306+
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),
307+
btagDeepFlavCvL = Var("?" + _mask + "?userFloat('btagDeepFlavCvLrefined'):max(bDiscriminator('pfDeepFlavourJetTags:probc')/(bDiscriminator('pfDeepFlavourJetTags:probc')+bDiscriminator('pfDeepFlavourJetTags:probuds')+bDiscriminator('pfDeepFlavourJetTags:probg')),-1)", float, precision=10),
308+
btagDeepFlavQG = Var("?" + _mask + "?userFloat('btagDeepFlavQGrefined'):max(bDiscriminator('pfDeepFlavourJetTags:probg')/(bDiscriminator('pfDeepFlavourJetTags:probg')+bDiscriminator('pfDeepFlavourJetTags:probuds')),-1)", float, precision=10),
309+
btagUParTAK4B = Var("?" + _mask + "?userFloat('btagUParTAK4Brefined'):max(bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:BvsAll'),-1)", float, precision=12),
310+
btagUParTAK4CvB = Var("?" + _mask + "?userFloat('btagUParTAK4CvBrefined'):max(bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:CvsB'),-1)", float, precision=12),
311+
btagUParTAK4CvL = Var("?" + _mask + "?userFloat('btagUParTAK4CvLrefined'):max(bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:CvsL'),-1)", float, precision=12),
312+
btagUParTAK4QvG = Var("?" + _mask + "?userFloat('btagUParTAK4QvGrefined'):max(bDiscriminator('pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:QvsG'),-1)", float, precision=12),
313+
)
314+
315+
# 4. Build pt_final as a userFloat on jets (mask applied), also do type 1 met correction
316+
# Note for future refinement: The mask definition above is also coded into the producer
317+
process.processRefinedJets = cms.EDProducer(
318+
"ProcessRefinedJets", # or "FastSimPuppiRefinedJetProducer" if that is your C++ class name
319+
jets = cms.InputTag("finalJetsPuppiWithRefined"),
320+
refinedPtName = cms.string("ptrefined"),
321+
maskBtagName = cms.string("pfUnifiedParticleTransformerAK4DiscriminatorsJetTags:BvsAll"),
322+
ptFinalName = cms.string("pt_final"),
323+
ptUnrefinedName = cms.string("pt_unrefined"),
324+
met = cms.InputTag("slimmedMETsPuppi"), # MET collection where type-1 refinement corrections
325+
)
326+
fastSim.toModify(
327+
process.jetPuppiTablesTask,
328+
process.jetPuppiTablesTask.add(process.processRefinedJets)
329+
)
330+
331+
#point PuppiMET Nano table to the *refined* MET collection
332+
fastSim.toModify(
333+
process.puppiMetTable,
334+
src = cms.InputTag("processRefinedJets", "Refined"),
335+
)
336+
337+
# And add _unrefined branches from the MET userFloats filled in C++
338+
fastSim.toModify(
339+
process.puppiMetTable.variables,
340+
pt_unrefined = Var("userFloat('pt_unrefined')", float, precision=10),
341+
phi_unrefined = Var("userFloat('phi_unrefined')", float, precision=10),
325342
)
343+
344+
# 5. Sort jets by pt_final only
345+
process.finalJetsPuppiSorted = cms.EDProducer(
346+
"JetSorter",
347+
src = cms.InputTag("processRefinedJets"),
348+
userFloatSorter = cms.string("pt_final"),
349+
descending = cms.bool(True),
350+
)
351+
fastSim.toModify(
352+
process.jetPuppiTablesTask,
353+
process.jetPuppiTablesTask.add(process.finalJetsPuppiSorted)
354+
)
355+
process.jetPuppiTable.src = cms.InputTag("finalJetsPuppiSorted")
356+
326357
return process
327358

359+
328360
# bind for the customise step
329361
nanoAOD_refineFastSim_puppiJet = nanoAOD_refineFastSim_puppiJet

0 commit comments

Comments
 (0)