@@ -220,22 +220,11 @@ def nanoAOD_addDeepInfoAK4(process,addParticleNet,addRobustParTAK4=False,addUnif
220220jetPuppiTablesTask = cms .Task (jetPuppiTable )
221221
222222from 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-
229223from 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
236225def 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
329361nanoAOD_refineFastSim_puppiJet = nanoAOD_refineFastSim_puppiJet
0 commit comments