ONNXRuntime/python/pfParticleNet_cff.py

0001 import FWCore.ParameterSet.Config as cms
0002
0003 from RecoBTag.FeatureTools.pfDeepBoostedJetTagInfos_cfi import pfDeepBoostedJetTagInfos
0004 from RecoBTag.ONNXRuntime.boostedJetONNXJetTagsProducer_cfi import boostedJetONNXJetTagsProducer
0005 from RecoBTag.ONNXRuntime.particleNetSonicJetTagsProducer_cfi import particleNetSonicJetTagsProducer as _particleNetSonicJetTagsProducer
0006 from RecoBTag.ONNXRuntime.pfParticleNetDiscriminatorsJetTags_cfi import pfParticleNetDiscriminatorsJetTags
0007 from RecoBTag.ONNXRuntime.pfMassDecorrelatedParticleNetDiscriminatorsJetTags_cfi import pfMassDecorrelatedParticleNetDiscriminatorsJetTags
0008 from Configuration.ProcessModifiers.particleNetSonicTriton_cff import particleNetSonicTriton
0009 from Configuration.ProcessModifiers.particleNetPTSonicTriton_cff import particleNetPTSonicTriton
0010
0011 pfParticleNetTagInfos = pfDeepBoostedJetTagInfos.clone(
0012     use_puppiP4 = False
0013 )
0014
0015 pfParticleNetJetTags = boostedJetONNXJetTagsProducer.clone(
0016     src = 'pfParticleNetTagInfos',
0017     preprocess_json = 'RecoBTag/Combined/data/ParticleNetAK8/General/V01/preprocess.json',
0018     model_path = 'RecoBTag/Combined/data/ParticleNetAK8/General/V01/modelfile/model.onnx',
0019     flav_names = ["probTbcq",  "probTbqq",  "probTbc",   "probTbq",  "probTbel", "probTbmu", "probTbta",
0020                   "probWcq",   "probWqq",   "probZbb",   "probZcc",  "probZqq",  "probHbb", "probHcc",
0021                   "probHqqqq", "probQCDbb", "probQCDcc", "probQCDb", "probQCDc", "probQCDothers"],
0022 )
0023
0024 particleNetSonicTriton.toReplaceWith(pfParticleNetJetTags, _particleNetSonicJetTagsProducer.clone(
0025     src = 'pfParticleNetTagInfos',
0026     preprocess_json = 'RecoBTag/Combined/data/ParticleNetAK8/General/V01/preprocess.json',
0027     Client = cms.PSet(
0028         timeout = cms.untracked.uint32(300),
0029         mode = cms.string("Async"),
0030         modelName = cms.string("particlenet"),
0031         modelConfigPath = cms.FileInPath("RecoBTag/Combined/data/models/particlenet/config.pbtxt"),
0032         modelVersion = cms.string(""),
0033         verbose = cms.untracked.bool(False),
0034         allowedTries = cms.untracked.uint32(0),
0035         useSharedMemory = cms.untracked.bool(True),
0036         compression = cms.untracked.string(""),
0037     ),
0038     flav_names = pfParticleNetJetTags.flav_names,
0039 ))
0040
0041 (particleNetSonicTriton & particleNetPTSonicTriton).toModify(pfParticleNetJetTags,
0042     preprocess_json = 'RecoBTag/Combined/data/ParticleNetAK8/General/V01/preprocess_PT.json',
0043     Client = dict(
0044         modelName = "particlenet_PT",
0045         modelConfigPath = cms.FileInPath("RecoBTag/Combined/data/models/particlenet_PT/config.pbtxt"),
0046     )
0047 )
0048
0049 pfMassDecorrelatedParticleNetJetTags = boostedJetONNXJetTagsProducer.clone(
0050     src = 'pfParticleNetTagInfos',
0051     preprocess_json = 'RecoBTag/Combined/data/ParticleNetAK8/MD-2prong/V01/preprocess.json',
0052     model_path = 'RecoBTag/Combined/data/ParticleNetAK8/MD-2prong/V01/modelfile/model.onnx',
0053     flav_names = ["probXbb", "probXcc", "probXqq", "probQCDbb", "probQCDcc",
0054                   "probQCDb", "probQCDc", "probQCDothers"],
0055 )
0056
0057 particleNetSonicTriton.toReplaceWith(pfMassDecorrelatedParticleNetJetTags, _particleNetSonicJetTagsProducer.clone(
0058     src = 'pfParticleNetTagInfos',
0059     preprocess_json = 'RecoBTag/Combined/data/ParticleNetAK8/MD-2prong/V01/preprocess.json',
0060     Client = cms.PSet(
0061         timeout = cms.untracked.uint32(300),
0062         modelName = cms.string("particlenet_AK8_MD-2prong"),
0063         mode = cms.string("Async"),
0064         modelConfigPath = cms.FileInPath("RecoBTag/Combined/data/models/particlenet_AK8_MD-2prong/config.pbtxt"),
0065         modelVersion = cms.string(""),
0066         verbose = cms.untracked.bool(False),
0067         allowedTries = cms.untracked.uint32(0),
0068     ),
0069     flav_names = pfMassDecorrelatedParticleNetJetTags.flav_names,
0070 ))
0071
0072 (particleNetSonicTriton & particleNetPTSonicTriton).toModify(pfMassDecorrelatedParticleNetJetTags,
0073     preprocess_json = 'RecoBTag/Combined/data/ParticleNetAK8/MD-2prong/V01/preprocess_PT.json',
0074     Client = dict(
0075         modelName = "particlenet_AK8_MD-2prong_PT",
0076         modelConfigPath = cms.FileInPath("RecoBTag/Combined/data/models/particlenet_AK8_MD-2prong_PT/config.pbtxt"),
0077     )
0078 )
0079
0080 pfParticleNetMassRegressionJetTags = boostedJetONNXJetTagsProducer.clone(
0081     src = 'pfParticleNetTagInfos',
0082     preprocess_json = 'RecoBTag/Combined/data/ParticleNetAK8/MassRegression/V01/preprocess.json',
0083     model_path = 'RecoBTag/Combined/data/ParticleNetAK8/MassRegression/V01/modelfile/model.onnx',
0084     flav_names = ["mass"],
0085 )
0086
0087 particleNetSonicTriton.toReplaceWith(pfParticleNetMassRegressionJetTags, _particleNetSonicJetTagsProducer.clone(
0088     src = 'pfParticleNetTagInfos',
0089     preprocess_json = 'RecoBTag/Combined/data/ParticleNetAK8/MassRegression/V01/preprocess.json',
0090     Client = cms.PSet(
0091         timeout = cms.untracked.uint32(300),
0092         modelName = cms.string("particlenet_AK8_MassRegression"),
0093         mode = cms.string("Async"),
0094         modelConfigPath = cms.FileInPath("RecoBTag/Combined/data/models/particlenet_AK8_MassRegression/config.pbtxt"),
0095         modelVersion = cms.string(""),
0096         verbose = cms.untracked.bool(False),
0097         allowedTries = cms.untracked.uint32(0),
0098     ),
0099     flav_names = pfParticleNetMassRegressionJetTags.flav_names,
0100 ))
0101
0102 (particleNetSonicTriton & particleNetPTSonicTriton).toModify(pfParticleNetMassRegressionJetTags,
0103     preprocess_json = 'RecoBTag/Combined/data/ParticleNetAK8/MassRegression/V01/preprocess_PT.json',
0104     Client = dict(
0105         modelName = "particlenet_AK8_MassRegression_PT",
0106         modelConfigPath = cms.FileInPath("RecoBTag/Combined/data/models/particlenet_AK8_MassRegression_PT/config.pbtxt"),
0107     )
0108 )
0109
0110 from CommonTools.PileupAlgos.Puppi_cff import puppi
0111 from CommonTools.RecoAlgos.primaryVertexAssociation_cfi import primaryVertexAssociation
0112
0113 # This task is not used, useful only if we run it from RECO jets (RECO/AOD)
0114 pfParticleNetTask = cms.Task(puppi, primaryVertexAssociation, pfParticleNetTagInfos,
0115                              pfParticleNetJetTags, pfMassDecorrelatedParticleNetJetTags, pfParticleNetMassRegressionJetTags,
0116                              pfParticleNetDiscriminatorsJetTags, pfMassDecorrelatedParticleNetDiscriminatorsJetTags)
0117
0118 # declare all the discriminators
0119 # nominal: probs
0120 _pfParticleNetJetTagsProbs = ['pfParticleNetJetTags:' + flav_name
0121                               for flav_name in pfParticleNetJetTags.flav_names]
0122
0123 # nominal: meta-taggers
0124 _pfParticleNetJetTagsMetaDiscrs = ['pfParticleNetDiscriminatorsJetTags:' + disc.name.value()
0125                                    for disc in pfParticleNetDiscriminatorsJetTags.discriminators]
0126 # mass-decorrelated: probs
0127 _pfMassDecorrelatedParticleNetJetTagsProbs = ['pfMassDecorrelatedParticleNetJetTags:' + flav_name
0128                               for flav_name in pfMassDecorrelatedParticleNetJetTags.flav_names]
0129
0130 # mass-decorrelated: meta-taggers
0131 _pfMassDecorrelatedParticleNetJetTagsMetaDiscrs = ['pfMassDecorrelatedParticleNetDiscriminatorsJetTags:' + disc.name.value()
0132                                    for disc in pfMassDecorrelatedParticleNetDiscriminatorsJetTags.discriminators]
0133
0134 _pfParticleNetMassRegressionOutputs = ['pfParticleNetMassRegressionJetTags:' + flav_name
0135                                        for flav_name in pfParticleNetMassRegressionJetTags.flav_names]
0136
0137 _pfParticleNetJetTagsAll = _pfParticleNetJetTagsProbs + _pfParticleNetJetTagsMetaDiscrs + \
0138     _pfMassDecorrelatedParticleNetJetTagsProbs + _pfMassDecorrelatedParticleNetJetTagsMetaDiscrs
0139
0140 _pfParticleNetMassCorrelatedJetTagsAll = _pfParticleNetJetTagsProbs + _pfParticleNetJetTagsMetaDiscrs