ONNXRuntime/test/test_particle_net_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/RunIIFall17MiniAODv2/TTToHadronic_TuneCP5_13TeV-powheg-pythia8/MINIAODSIM/PU2017_12Apr2018_94X_mc2017_realistic_v14-v1/90000/DCFE3F5F-AE42-E811-B6DB-008CFAF72A64.root'
0007 options.inputFiles = 'file:/storage/local/data1/home/jduarte1/forPatrick/FFA0194D-1BBC-EF4F-9B8F-8FBED2C62FC8.root'
0008 #options.inputFiles = '/store/mc/RunIISummer19UL17MiniAOD/TTToHadronic_TuneCP5_13TeV-powheg-pythia8/MINIAODSIM/106X_mc2017_realistic_v6-v4/30000/FFA0194D-1BBC-EF4F-9B8F-8FBED2C62FC8.root'
0009 #options.inputFiles = 'file:FFA0194D-1BBC-EF4F-9B8F-8FBED2C62FC8.root'
0010 options.maxEvents = 100
0011 options.parseArguments()
0012
0013 from Configuration.ProcessModifiers.enableSonicTriton_cff import enableSonicTriton
0014 process = cms.Process('PATtest',enableSonicTriton)
0015
0016 ## MessageLogger
0017 process.load("FWCore.MessageLogger.MessageLogger_cfi")
0018 process.MessageLogger.cerr.FwkReport.reportEvery = 1
0019
0020
0021 ## Options and Output Report
0022 process.options = cms.untracked.PSet( wantSummary = cms.untracked.bool(True) )
0023
0024 ## Source
0025 process.source = cms.Source("PoolSource",
0026     fileNames=cms.untracked.vstring(options.inputFiles)
0027 )
0028 ## Maximal Number of Events
0029 process.maxEvents = cms.untracked.PSet(input=cms.untracked.int32(options.maxEvents))
0030
0031 process.load("HeterogeneousCore.SonicTriton.TritonService_cff")
0032 process.TritonService.verbose = True
0033 # fallback server
0034 process.TritonService.fallback.enable = True
0035 process.TritonService.fallback.verbose = True
0036 process.TritonService.servers.append(
0037     cms.PSet(
0038         name = cms.untracked.string("default"),
0039         #address = cms.untracked.string("prp-gpu-1.t2.ucsd.edu"),
0040         address = cms.untracked.string("ailab01.fnal.gov"),
0041         port = cms.untracked.uint32(8001),
0042     )
0043 )
0044
0045
0046 ## Geometry and Detector Conditions (needed for a few patTuple production steps)
0047 process.load("Configuration.Geometry.GeometryRecoDB_cff")
0048 process.load("Configuration.StandardSequences.FrontierConditions_GlobalTag_cff")
0049 from Configuration.AlCa.GlobalTag import GlobalTag
0050 process.GlobalTag = GlobalTag(process.GlobalTag, 'auto:phase1_2017_realistic')
0051 process.load("Configuration.StandardSequences.MagneticField_cff")
0052
0053 ## Output Module Configuration (expects a path 'p')
0054 from PhysicsTools.PatAlgos.patEventContent_cff import patEventContentNoCleaning
0055 process.out = cms.OutputModule("PoolOutputModule",
0056                                fileName = cms.untracked.string('patTuple.root'),
0057                                ## save only events passing the full path
0058                                #SelectEvents = cms.untracked.PSet( SelectEvents = cms.vstring('p') ),
0059                                ## save PAT output; you need a '*' to unpack the list of commands
0060                                ## 'patEventContent'
0061                                outputCommands = cms.untracked.vstring('drop *', *patEventContentNoCleaning )
0062                                )
0063
0064 patAlgosToolsTask = getPatAlgosToolsTask(process)
0065 process.outpath = cms.EndPath(process.out, patAlgosToolsTask)
0066
0067 ## and add them to the event content
0068 from PhysicsTools.PatAlgos.tools.jetTools import updateJetCollection
0069 from RecoBTag.ONNXRuntime.pfParticleNet_cff import _pfParticleNetJetTagsProbs as pfParticleNetJetTagsProbs
0070 from RecoBTag.ONNXRuntime.pfParticleNet_cff import _pfParticleNetSonicJetTagsProbs as pfParticleNetSonicJetTagsProbs
0071
0072 updateJetCollection(
0073    process,
0074    jetSource = cms.InputTag('slimmedJetsAK8'),
0075    pvSource = cms.InputTag('offlineSlimmedPrimaryVertices'),
0076    svSource = cms.InputTag('slimmedSecondaryVertices'),
0077    rParam = 0.8,
0078    jetCorrections = ('AK8PFPuppi', cms.vstring(['L2Relative', 'L3Absolute']), 'None'),
0079    btagDiscriminators = pfParticleNetJetTagsProbs+pfParticleNetSonicJetTagsProbs
0080    )
0081
0082 from Configuration.EventContent.EventContent_cff import MINIAODSIMEventContent
0083 process.out.outputCommands.append('keep *_slimmedJetsAK8*_*_*')
0084 process.out.outputCommands.append('keep *_offlineSlimmedPrimaryVertices*_*_*')
0085 process.out.outputCommands.append('keep *_slimmedSecondaryVertices*_*_*')
0086 process.out.outputCommands.append('keep *_selectedPatJets*_*_*')
0087 process.out.outputCommands.append('keep *_selectedUpdatedPatJets*_*_*')
0088 process.out.outputCommands.append('keep *_updatedPatJets*_*_*')
0089
0090 process.out.fileName = 'test_particle_net_MINIAODSIM_noragged.root'
0091
0092