Skip to content

Commit

Permalink
Adding changes to config files and producer to enable SONIC for new v…
Browse files Browse the repository at this point in the history 
…ersions of particlenet and higgs interaction network
  • Loading branch information
@wpmccormack
wpmccormack committed Oct 27, 2023
1 parent e4bfeac commit 6855d27
Show file tree
Hide file tree
Showing 5 changed files with 192 additions and 11 deletions.
92 changes: 81 additions & 11 deletions RecoBTag/ONNXRuntime/plugins/ParticleNetSonicJetTagsProducer.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,14 @@ class ParticleNetSonicJetTagsProducer : public TritonEDProducer<> {
std::unordered_map<std::string, PreprocessParams> prep_info_map_; // preprocessing info for each input group
bool debug_ = false;
bool skippedInference_ = false;
constexpr static unsigned numParticleGroups_ = 3;
std::vector<bool> emptyJets_;
unsigned numParticleGroups_ = 0;
unsigned numVertexGroups_ = 0;
unsigned numLostTrackGroups_ = 0;
bool countedInputs = false;
std::string particleNameExample_;
std::string vertexNameExample_;
std::string losttrackNameExample_;
};

ParticleNetSonicJetTagsProducer::ParticleNetSonicJetTagsProducer(const edm::ParameterSet &iConfig)
Expand Down Expand Up @@ -86,6 +93,8 @@ ParticleNetSonicJetTagsProducer::ParticleNetSonicJetTagsProducer(const edm::Para
for (const auto &flav_name : flav_names_) {
produces<JetTagCollection>(flav_name);
}

emptyJets_.clear();
}

ParticleNetSonicJetTagsProducer::~ParticleNetSonicJetTagsProducer() {}
Expand Down Expand Up @@ -131,26 +140,74 @@ void ParticleNetSonicJetTagsProducer::acquire(edm::Event const &iEvent, edm::Eve
iEvent.getByToken(src_, tag_infos);
client_->setBatchSize(tag_infos->size());
skippedInference_ = false;

emptyJets_.clear();

if (!countedInputs) {
int model_input_size = input_names_.size();
for (int n = 0; n < model_input_size; n++) {
if (prep_info_map_.at(input_names_.at(n)).var_names.at(0).find("pf") != std::string::npos) {
if (numParticleGroups_ == 0) {
particleNameExample_ = prep_info_map_.at(input_names_.at(n)).var_names.at(0);
}
numParticleGroups_++;
} else if (prep_info_map_.at(input_names_.at(n)).var_names.at(0).find("sv") != std::string::npos) {
if (numVertexGroups_ == 0) {
vertexNameExample_ = prep_info_map_.at(input_names_.at(n)).var_names.at(0);
}
numVertexGroups_++;
} else if (prep_info_map_.at(input_names_.at(n)).var_names.at(0).find("lt") != std::string::npos) {
if (numLostTrackGroups_ == 0) {
losttrackNameExample_ = prep_info_map_.at(input_names_.at(n)).var_names.at(0);
}
numLostTrackGroups_++;
}
}

countedInputs = true;
}

if (!tag_infos->empty()) {
unsigned int maxParticles = 0;
unsigned int maxVertices = 0;
unsigned int maxLT = 0;
unsigned int numParticles;
unsigned int numVertices;
for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) {
maxParticles = std::max(maxParticles,
static_cast<unsigned int>(((*tag_infos)[jet_n]).features().get("pfcand_etarel").size()));
maxVertices =
std::max(maxVertices, static_cast<unsigned int>(((*tag_infos)[jet_n]).features().get("sv_etarel").size()));
numParticles = static_cast<unsigned int>(((*tag_infos)[jet_n]).features().get(particleNameExample_).size());
numVertices = static_cast<unsigned int>(((*tag_infos)[jet_n]).features().get(vertexNameExample_).size());
maxParticles = std::max(maxParticles, numParticles);
maxVertices = std::max(maxVertices, numVertices);

if (numParticles == 0 && numVertices == 0) {
emptyJets_.push_back(true);
} else {
emptyJets_.push_back(false);
}

if (!(losttrackNameExample_.empty()) && numParticles > 0) {
maxLT = std::max(maxLT,
static_cast<unsigned int>(((*tag_infos)[jet_n]).features().get(losttrackNameExample_).size()));
}
}
if (maxParticles == 0 && maxVertices == 0) {

if (maxParticles == 0 && maxVertices == 0 && maxLT == 0) {
client_->setBatchSize(0);
skippedInference_ = true;
return;
}

unsigned int minPartFromJSON = prep_info_map_.at(input_names_[0]).min_length;
unsigned int maxPartFromJSON = prep_info_map_.at(input_names_[0]).max_length;
unsigned int minVertFromJSON = prep_info_map_.at(input_names_[3]).min_length;
unsigned int maxVertFromJSON = prep_info_map_.at(input_names_[3]).max_length;
unsigned int minVertFromJSON = prep_info_map_.at(input_names_[numParticleGroups_]).min_length;
unsigned int maxVertFromJSON = prep_info_map_.at(input_names_[numParticleGroups_]).max_length;
maxParticles = std::clamp(maxParticles, minPartFromJSON, maxPartFromJSON);
maxVertices = std::clamp(maxVertices, minVertFromJSON, maxVertFromJSON);
if (!(losttrackNameExample_.empty())) {
unsigned int minLTFromJSON = prep_info_map_.at(input_names_[numParticleGroups_ + numVertexGroups_]).min_length;
unsigned int maxLTFromJSON = prep_info_map_.at(input_names_[numParticleGroups_ + numVertexGroups_]).max_length;
maxLT = std::clamp(maxLT, minLTFromJSON, maxLTFromJSON);
}

for (unsigned igroup = 0; igroup < input_names_.size(); ++igroup) {
const auto &group_name = input_names_[igroup];
Expand All @@ -159,9 +216,12 @@ void ParticleNetSonicJetTagsProducer::acquire(edm::Event const &iEvent, edm::Eve
if (igroup < numParticleGroups_) {
input.setShape(1, maxParticles);
target = maxParticles;
} else {
} else if (igroup < (numParticleGroups_ + numVertexGroups_)) {
input.setShape(1, maxVertices);
target = maxVertices;
} else {
input.setShape(1, maxLT);
target = maxLT;
}
auto tdata = input.allocate<float>(true);
for (unsigned jet_n = 0; jet_n < tag_infos->size(); ++jet_n) {
Expand All @@ -172,7 +232,13 @@ void ParticleNetSonicJetTagsProducer::acquire(edm::Event const &iEvent, edm::Eve
// transform/pad
for (unsigned i = 0; i < prep_params.var_names.size(); ++i) {
const auto &varname = prep_params.var_names[i];
const auto &raw_value = taginfo.features().get(varname);
std::vector<float> bare(0);
std::vector<float> raw_value;
if (!emptyJets_.at(jet_n)) {
raw_value = taginfo.features().get(varname);
} else {
raw_value = bare;
}
const auto &info = prep_params.info(varname);
int insize = center_norm_pad_halfRagged(raw_value,
info.center,
Expand Down Expand Up @@ -234,7 +300,11 @@ void ParticleNetSonicJetTagsProducer::produce(edm::Event &iEvent,
const auto &taginfo = (*tag_infos)[jet_n];
const auto &jet_ref = tag_infos->at(jet_n).jet();

if (!taginfo.features().empty()) {
if (emptyJets_.at(jet_n)) {
for (std::size_t flav_n = 0; flav_n < flav_names_.size(); flav_n++) {
(*(output_tags[flav_n]))[jet_ref] = 0.;
}
} else if (!taginfo.features().empty()) {
for (std::size_t flav_n = 0; flav_n < flav_names_.size(); flav_n++) {
(*(output_tags[flav_n]))[jet_ref] = outputs_from_server[jet_n][flav_n];
}
Expand Down
20 changes: 20 additions & 0 deletions RecoBTag/ONNXRuntime/python/pfHiggsInteractionNet_cff.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,9 @@
from RecoBTag.FeatureTools.pfDeepBoostedJetTagInfos_cfi import pfDeepBoostedJetTagInfos as _pfDeepBoostedJetTagInfos
from RecoBTag.ONNXRuntime.boostedJetONNXJetTagsProducer_cfi import boostedJetONNXJetTagsProducer
from RecoBTag.ONNXRuntime.Parameters.HiggsInteractionNet.V00.pfHiggsInteractionNetPreprocessParams_cfi import pfHiggsInteractionNetPreprocessParams
from RecoBTag.ONNXRuntime.particleNetSonicJetTagsProducer_cfi import particleNetSonicJetTagsProducer as _particleNetSonicJetTagsProducer
from Configuration.ProcessModifiers.particleNetSonicTriton_cff import particleNetSonicTriton
from Configuration.ProcessModifiers.particleNetPTSonicTriton_cff import particleNetPTSonicTriton

# modify default parameters for tag infos
pfHiggsInteractionNetTagInfos = _pfDeepBoostedJetTagInfos.clone(
Expand All @@ -21,6 +24,23 @@
flav_names = [ 'probQCD', 'probHbb' ]
)

particleNetSonicTriton.toReplaceWith(pfHiggsInteractionNetTags, _particleNetSonicJetTagsProducer.clone(
src = 'pfHiggsInteractionNetTagInfos',
preprocess_json = 'RecoBTag/Combined/data/models/higgsInteractionNet/preprocess.json',
Client = cms.PSet(
timeout = cms.untracked.uint32(300),
mode = cms.string("Async"),
modelName = cms.string("higgsInteractionNet"),
modelConfigPath = cms.FileInPath("RecoBTag/Combined/data/models/higgsInteractionNet/config.pbtxt"),
modelVersion = cms.string(""),
verbose = cms.untracked.bool(False),
allowedTries = cms.untracked.uint32(0),
useSharedMemory = cms.untracked.bool(True),
compression = cms.untracked.string(""),
),
flav_names = pfHiggsInteractionNetTags.flav_names,
))

# declare all the discriminators
# nominal: probs
_pfHiggsInteractionNetTagsProbs = ['pfHiggsInteractionNetTags:' + flav_name
Expand Down
71 changes: 71 additions & 0 deletions RecoBTag/ONNXRuntime/python/pfParticleNetFromMiniAODAK4_cff.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,9 @@
from RecoBTag.ONNXRuntime.boostedJetONNXJetTagsProducer_cfi import boostedJetONNXJetTagsProducer
from RecoBTag.FeatureTools.ParticleNetFeatureEvaluator_cfi import ParticleNetFeatureEvaluator
from RecoBTag.ONNXRuntime.pfParticleNetFromMiniAODAK4DiscriminatorsJetTags_cfi import *
from RecoBTag.ONNXRuntime.particleNetSonicJetTagsProducer_cfi import particleNetSonicJetTagsProducer as _particleNetSonicJetTagsProducer
from Configuration.ProcessModifiers.particleNetSonicTriton_cff import particleNetSonicTriton
from Configuration.ProcessModifiers.particleNetPTSonicTriton_cff import particleNetPTSonicTriton

pfParticleNetFromMiniAODAK4CHSCentralTagInfos = ParticleNetFeatureEvaluator.clone(
jets = "slimmedJets",
Expand Down Expand Up @@ -45,27 +48,95 @@
flav_names = ['probmu','probele','probtaup1h0p','probtaup1h1p','probtaup1h2p','probtaup3h0p','probtaup3h1p','probtaum1h0p','probtaum1h1p','probtaum1h2p','probtaum3h0p','probtaum3h1p','probb','probc','probuds','probg','ptcorr','ptreshigh','ptreslow','ptnu'],
)

particleNetSonicTriton.toReplaceWith(pfParticleNetFromMiniAODAK4CHSCentralJetTags, _particleNetSonicJetTagsProducer.clone(
src = 'pfParticleNetFromMiniAODAK4CHSCentralTagInfos',
preprocess_json = 'RecoBTag/Combined/data/ParticleNetFromMiniAODAK4/CHS/Central/preprocess.json',
Client = cms.PSet(
timeout = cms.untracked.uint32(300),
mode = cms.string("Async"),
modelName = cms.string("particleNetFromMiniAODAK4CHSCentral"),
modelConfigPath = cms.FileInPath("RecoBTag/Combined/data/models/particleNetFromMiniAODAK4CHSCentral/config.pbtxt"),
modelVersion = cms.string(""),
verbose = cms.untracked.bool(False),
allowedTries = cms.untracked.uint32(0),
useSharedMemory = cms.untracked.bool(True),
compression = cms.untracked.string(""),
),
flav_names = pfParticleNetFromMiniAODAK4CHSCentralJetTags.flav_names,
))

pfParticleNetFromMiniAODAK4CHSForwardJetTags = boostedJetONNXJetTagsProducer.clone(
src = 'pfParticleNetFromMiniAODAK4CHSForwardTagInfos',
preprocess_json = 'RecoBTag/Combined/data/ParticleNetFromMiniAODAK4/CHS/Forward/preprocess.json',
model_path = 'RecoBTag/Combined/data/ParticleNetFromMiniAODAK4/CHS/Forward/particle-net.onnx',
flav_names = ['probq','probg','ptcorr','ptreshigh','ptreslow','ptnu'],
)

particleNetSonicTriton.toReplaceWith(pfParticleNetFromMiniAODAK4CHSForwardJetTags, _particleNetSonicJetTagsProducer.clone(
src = 'pfParticleNetFromMiniAODAK4CHSForwardTagInfos',
preprocess_json = 'RecoBTag/Combined/data/ParticleNetFromMiniAODAK4/CHS/Central/preprocess.json',
Client = cms.PSet(
timeout = cms.untracked.uint32(300),
mode = cms.string("Async"),
modelName = cms.string("particleNetFromMiniAODAK4CHSForward"),
modelConfigPath = cms.FileInPath("RecoBTag/Combined/data/models/particleNetFromMiniAODAK4CHSForward/config.pbtxt"),
modelVersion = cms.string(""),
verbose = cms.untracked.bool(False),
allowedTries = cms.untracked.uint32(0),
useSharedMemory = cms.untracked.bool(True),
compression = cms.untracked.string(""),
),
flav_names = pfParticleNetFromMiniAODAK4CHSForwardJetTags.flav_names,
))

pfParticleNetFromMiniAODAK4PuppiCentralJetTags = boostedJetONNXJetTagsProducer.clone(
src = 'pfParticleNetFromMiniAODAK4PuppiCentralTagInfos',
preprocess_json = 'RecoBTag/Combined/data/ParticleNetFromMiniAODAK4/PUPPI/Central/preprocess.json',
model_path = 'RecoBTag/Combined/data/ParticleNetFromMiniAODAK4/PUPPI/Central/particle-net.onnx',
flav_names = ['probmu','probele','probtaup1h0p','probtaup1h1p','probtaup1h2p','probtaup3h0p','probtaup3h1p','probtaum1h0p','probtaum1h1p','probtaum1h2p','probtaum3h0p','probtaum3h1p','probb','probc','probuds','probg','ptcorr','ptreshigh','ptreslow','ptnu'],
)

particleNetSonicTriton.toReplaceWith(pfParticleNetFromMiniAODAK4PuppiCentralJetTags, _particleNetSonicJetTagsProducer.clone(
src = 'pfParticleNetFromMiniAODAK4PuppiCentralTagInfos',
preprocess_json = 'RecoBTag/Combined/data/ParticleNetFromMiniAODAK4/PUPPI/Central/preprocess.json',
Client = cms.PSet(
timeout = cms.untracked.uint32(300),
mode = cms.string("Async"),
modelName = cms.string("particleNetFromMiniAODAK4PuppiCentral"),
modelConfigPath = cms.FileInPath("RecoBTag/Combined/data/models/particleNetFromMiniAODAK4PuppiCentral/config.pbtxt"),
modelVersion = cms.string(""),
verbose = cms.untracked.bool(False),
allowedTries = cms.untracked.uint32(0),
useSharedMemory = cms.untracked.bool(True),
compression = cms.untracked.string(""),
),
flav_names = pfParticleNetFromMiniAODAK4PuppiCentralJetTags.flav_names,
))

pfParticleNetFromMiniAODAK4PuppiForwardJetTags = boostedJetONNXJetTagsProducer.clone(
src = 'pfParticleNetFromMiniAODAK4PuppiForwardTagInfos',
preprocess_json = 'RecoBTag/Combined/data/ParticleNetFromMiniAODAK4/PUPPI/Forward/preprocess.json',
model_path = 'RecoBTag/Combined/data/ParticleNetFromMiniAODAK4/PUPPI/Forward/particle-net.onnx',
flav_names = ['probq','probg','ptcorr','ptreshigh','ptreslow','ptnu'],
)

particleNetSonicTriton.toReplaceWith(pfParticleNetFromMiniAODAK4PuppiForwardJetTags, _particleNetSonicJetTagsProducer.clone(
src = 'pfParticleNetFromMiniAODAK4PuppiForwardTagInfos',
preprocess_json = 'RecoBTag/Combined/data/ParticleNetFromMiniAODAK4/PUPPI/Forward/preprocess.json',
Client = cms.PSet(
timeout = cms.untracked.uint32(300),
mode = cms.string("Async"),
modelName = cms.string("particleNetFromMiniAODAK4PuppiForward"),
modelConfigPath = cms.FileInPath("RecoBTag/Combined/data/models/particleNetFromMiniAODAK4PuppiForward/config.pbtxt"),
modelVersion = cms.string(""),
verbose = cms.untracked.bool(False),
allowedTries = cms.untracked.uint32(0),
useSharedMemory = cms.untracked.bool(True),
compression = cms.untracked.string(""),
),
flav_names = pfParticleNetFromMiniAODAK4PuppiForwardJetTags.flav_names,
))

pfParticleNetFromMiniAODAK4CHSTask = cms.Task( pfParticleNetFromMiniAODAK4CHSCentralTagInfos, pfParticleNetFromMiniAODAK4CHSForwardTagInfos, pfParticleNetFromMiniAODAK4CHSCentralJetTags, pfParticleNetFromMiniAODAK4CHSForwardJetTags)
pfParticleNetFromMiniAODAK4PuppiTask = cms.Task( pfParticleNetFromMiniAODAK4PuppiCentralTagInfos, pfParticleNetFromMiniAODAK4PuppiForwardTagInfos, pfParticleNetFromMiniAODAK4PuppiCentralJetTags, pfParticleNetFromMiniAODAK4PuppiForwardJetTags)

Expand Down
19 changes: 19 additions & 0 deletions RecoBTag/ONNXRuntime/python/pfParticleNetFromMiniAODAK8_cff.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,9 @@
from RecoBTag.ONNXRuntime.boostedJetONNXJetTagsProducer_cfi import boostedJetONNXJetTagsProducer
from RecoBTag.FeatureTools.ParticleNetFeatureEvaluator_cfi import ParticleNetFeatureEvaluator
from RecoBTag.ONNXRuntime.pfParticleNetFromMiniAODAK8DiscriminatorsJetTags_cfi import pfParticleNetFromMiniAODAK8DiscriminatorsJetTags
from RecoBTag.ONNXRuntime.particleNetSonicJetTagsProducer_cfi import particleNetSonicJetTagsProducer as _particleNetSonicJetTagsProducer
from Configuration.ProcessModifiers.particleNetSonicTriton_cff import particleNetSonicTriton
from Configuration.ProcessModifiers.particleNetPTSonicTriton_cff import particleNetPTSonicTriton

pfParticleNetFromMiniAODAK8TagInfos = ParticleNetFeatureEvaluator.clone(
jets = "slimmedJetsAK8",
Expand All @@ -21,6 +24,22 @@
flav_names = ['probHtt','probHtm','probHte','probHbb', 'probHcc', 'probHqq', 'probHgg','probQCD2hf','probQCD1hf','probQCD0hf','masscorr'],
)

particleNetSonicTriton.toReplaceWith(pfParticleNetFromMiniAODAK8JetTags, _particleNetSonicJetTagsProducer.clone(
src = 'pfParticleNetFromMiniAODAK8TagInfos',
preprocess_json = 'RecoBTag/Combined/data/ParticleNetFromMiniAODAK8/preprocess.json',
Client = cms.PSet(
timeout = cms.untracked.uint32(300),
mode = cms.string("Async"),
modelName = cms.string("particleNetFromMiniAODAK8"),
modelConfigPath = cms.FileInPath("RecoBTag/Combined/data/models/particleNetFromMiniAODAK8/config.pbtxt"),
modelVersion = cms.string(""),
verbose = cms.untracked.bool(False),
allowedTries = cms.untracked.uint32(0),
useSharedMemory = cms.untracked.bool(True),
compression = cms.untracked.string(""),
),
flav_names = pfParticleNetFromMiniAODAK8JetTags.flav_names,
))

pfParticleNetFromMiniAODAK8Task = cms.Task( pfParticleNetFromMiniAODAK8TagInfos, pfParticleNetFromMiniAODAK8JetTags)

Expand Down
1 change: 1 addition & 0 deletions RecoBTag/ONNXRuntime/python/pfParticleNet_cff.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,6 +46,7 @@
)
)


pfMassDecorrelatedParticleNetJetTags = boostedJetONNXJetTagsProducer.clone(
src = 'pfParticleNetTagInfos',
preprocess_json = 'RecoBTag/Combined/data/ParticleNetAK8/MD-2prong/V01/preprocess.json',
Expand Down

0 comments on commit 6855d27

Please sign in to comment.