ONNXRuntime/test/test_particle_net_ak4_cfg.py

0001 import FWCore.ParameterSet.Config as cms
0002 from PhysicsTools.PatAlgos.tools.helpers import getPatAlgosToolsTask
0003
0004 from FWCore.ParameterSet.VarParsing import VarParsing
0005 options = VarParsing('analysis')
0006 #options.inputFiles = '/store/mc/RunIISummer19UL17MiniAOD/TTToHadronic_TuneCP5_13TeV-powheg-pythia8/MINIAODSIM/106X_mc2017_realistic_v6-v4/30000/FFA0194D-1BBC-EF4F-9B8F-8FBED2C62FC8.root'
0007 options.inputFiles = 'file:/storage/local/data1/home/jduarte1/forPatrick/FFA0194D-1BBC-EF4F-9B8F-8FBED2C62FC8.root'
0008 options.maxEvents = 1000
0009 options.parseArguments()
0010
0011 from Configuration.ProcessModifiers.enableSonicTriton_cff import enableSonicTriton
0012 process = cms.Process('PATtest',enableSonicTriton)
0013
0014 process = cms.Process("PATtest")
0015
0016 ## MessageLogger
0017 process.load("FWCore.MessageLogger.MessageLogger_cfi")
0018 process.MessageLogger.cerr.FwkReport.reportEvery = 1
0019
0020 keepMsgs = ['TritonClient','TritonService']
0021 keepMsgs.extend(['BoostedJetONNXJetTagsProducer'])
0022 keepMsgs.extend(['ParticleNetSonicJetTagsProducer', 'ParticleNetSonicJetTagsProducer:TritonClient'])
0023 for msg in keepMsgs:
0024     setattr(process.MessageLogger.cerr,msg,
0025         cms.untracked.PSet(
0026             limit = cms.untracked.int32(10000000),
0027         )
0028     )
0029
0030
0031 ## Options and Output Report
0032 process.options = cms.untracked.PSet( wantSummary = cms.untracked.bool(True) )
0033
0034 ## Source
0035 process.source = cms.Source("PoolSource",
0036     fileNames=cms.untracked.vstring(options.inputFiles)
0037 )
0038 ## Maximal Number of Events
0039 process.maxEvents = cms.untracked.PSet(input=cms.untracked.int32(options.maxEvents))
0040
0041 process.load("HeterogeneousCore.SonicTriton.TritonService_cff")
0042 process.TritonService.verbose = True
0043 # fallback server
0044 process.TritonService.fallback.enable = True
0045 process.TritonService.fallback.verbose = True
0046 process.TritonService.fallback.useGPU = False
0047 process.TritonService.servers.append(
0048     cms.PSet(
0049         name = cms.untracked.string("default"),
0050         #address = cms.untracked.string("prp-gpu-1.t2.ucsd.edu"),
0051         address = cms.untracked.string("ailab01.fnal.gov"),
0052         port = cms.untracked.uint32(8001),
0053     )
0054 )
0055
0056
0057 ## Geometry and Detector Conditions (needed for a few patTuple production steps)
0058 process.load("Configuration.Geometry.GeometryRecoDB_cff")
0059 process.load("Configuration.StandardSequences.FrontierConditions_GlobalTag_cff")
0060 from Configuration.AlCa.GlobalTag import GlobalTag
0061 process.GlobalTag = GlobalTag(process.GlobalTag, 'auto:phase1_2017_realistic')
0062 process.load("Configuration.StandardSequences.MagneticField_cff")
0063
0064 ## Output Module Configuration (expects a path 'p')
0065 from PhysicsTools.PatAlgos.patEventContent_cff import patEventContentNoCleaning
0066 process.out = cms.OutputModule("PoolOutputModule",
0067                                fileName = cms.untracked.string('patTuple.root'),
0068                                ## save only events passing the full path
0069                                #SelectEvents = cms.untracked.PSet( SelectEvents = cms.vstring('p') ),
0070                                ## save PAT output; you need a '*' to unpack the list of commands
0071                                ## 'patEventContent'
0072                                outputCommands = cms.untracked.vstring('drop *', *patEventContentNoCleaning )
0073                                )
0074
0075 patAlgosToolsTask = getPatAlgosToolsTask(process)
0076 process.outpath = cms.EndPath(process.out, patAlgosToolsTask)
0077
0078 ## and add them to the event content
0079 from PhysicsTools.PatAlgos.tools.jetTools import updateJetCollection
0080 from RecoBTag.ONNXRuntime.pfParticleNetAK4_cff import _pfParticleNetAK4JetTagsAll as pfParticleNetAK4JetTagsAll
0081
0082 updateJetCollection(
0083    process,
0084    jetSource = cms.InputTag('slimmedJets'),
0085    pvSource = cms.InputTag('offlineSlimmedPrimaryVertices'),
0086    svSource = cms.InputTag('slimmedSecondaryVertices'),
0087    jetCorrections = ('AK4PFchs', cms.vstring(['L1FastJet', 'L2Relative', 'L3Absolute']), 'None'),
0088    btagDiscriminators = [
0089           'pfDeepFlavourJetTags:probb',
0090           'pfDeepFlavourJetTags:probbb',
0091           'pfDeepFlavourJetTags:problepb',
0092           'pfDeepFlavourJetTags:probc',
0093           'pfDeepFlavourJetTags:probuds',
0094           'pfDeepFlavourJetTags:probg',
0095        ] + pfParticleNetAK4JetTagsAll
0096    )
0097
0098 from Configuration.EventContent.EventContent_cff import MINIAODSIMEventContent
0099 process.out.outputCommands.append('keep *_slimmedJets*_*_*')
0100 process.out.outputCommands.append('keep *_offlineSlimmedPrimaryVertices*_*_*')
0101 process.out.outputCommands.append('keep *_slimmedSecondaryVertices*_*_*')
0102 process.out.outputCommands.append('keep *_selectedPatJets*_*_*')
0103 process.out.outputCommands.append('keep *_selectedUpdatedPatJets*_*_*')
0104 process.out.outputCommands.append('keep *_updatedPatJets*_*_*')
0105
0106 process.out.fileName = 'test_particle_net_ak4_MINIAODSIM.root'