Skip to content

Commit

Permalink
Add use of unsubtracted jets in b-tagger algorithms
Browse files Browse the repository at this point in the history
  • Loading branch information
Andre Govinda Stahl Leiton committed Mar 15, 2024
1 parent d389c1c commit 4e1f0e4
Show file tree
Hide file tree
Showing 6 changed files with 191 additions and 13 deletions.
27 changes: 22 additions & 5 deletions RecoBTag/FeatureTools/plugins/DeepBoostedJetTagInfoProducer.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
#include "DataFormats/VertexReco/interface/VertexFwd.h"

#include "DataFormats/BTauReco/interface/DeepBoostedJetTagInfo.h"
#include "DataFormats/Common/interface/AssociationMap.h"

using namespace btagbtvdeep;

Expand All @@ -39,12 +40,13 @@ class DeepBoostedJetTagInfoProducer : public edm::stream::EDProducer<> {
typedef reco::VertexCompositePtrCandidateCollection SVCollection;
typedef reco::VertexCollection VertexCollection;
typedef edm::View<reco::Candidate> CandidateView;
typedef edm::AssociationMap<edm::OneToOne<reco::JetView, reco::JetView>> JetMatchMap;

void beginStream(edm::StreamID) override {}
void produce(edm::Event &, const edm::EventSetup &) override;
void endStream() override {}

void fillParticleFeatures(DeepBoostedJetFeatures &fts, const reco::Jet &jet);
void fillParticleFeatures(DeepBoostedJetFeatures &fts, const reco::Jet &unsubJet, const reco::Jet &jet);
void fillSVFeatures(DeepBoostedJetFeatures &fts, const reco::Jet &jet);
void fillParticleFeaturesHLT(DeepBoostedJetFeatures &fts, const reco::Jet &jet, const reco::VertexRefProd &PVRefProd);
void fillSVFeaturesHLT(DeepBoostedJetFeatures &fts, const reco::Jet &jet);
Expand All @@ -66,13 +68,15 @@ class DeepBoostedJetTagInfoProducer : public edm::stream::EDProducer<> {
const bool use_hlt_features_;

edm::EDGetTokenT<edm::View<reco::Jet>> jet_token_;
edm::EDGetTokenT<JetMatchMap> unsubjet_map_token_;
edm::EDGetTokenT<VertexCollection> vtx_token_;
edm::EDGetTokenT<SVCollection> sv_token_;
edm::EDGetTokenT<CandidateView> pfcand_token_;

bool use_puppi_value_map_;
bool use_pvasq_value_map_;
bool is_packed_pf_candidate_collection_;
bool use_unsubjet_map_;

edm::EDGetTokenT<edm::ValueMap<float>> puppi_value_map_token_;
edm::EDGetTokenT<edm::ValueMap<int>> pvasq_value_map_token_;
Expand Down Expand Up @@ -229,6 +233,7 @@ DeepBoostedJetTagInfoProducer::DeepBoostedJetTagInfoProducer(const edm::Paramete
pfcand_token_(consumes<CandidateView>(iConfig.getParameter<edm::InputTag>("pf_candidates"))),
use_puppi_value_map_(false),
use_pvasq_value_map_(false),
use_unsubjet_map_(false),
track_builder_token_(
esConsumes<TransientTrackBuilder, TransientTrackRecord>(edm::ESInputTag("", "TransientTrackBuilder"))),
covarianceVersion_(iConfig.getParameter<int>("covarianceVersion")),
Expand Down Expand Up @@ -300,6 +305,12 @@ DeepBoostedJetTagInfoProducer::DeepBoostedJetTagInfoProducer(const edm::Paramete
trkPhi_value_map_token_ = consumes<edm::ValueMap<float>>(trkPhi_value_map_tag);
}

const auto &unsubjet_map_tag = iConfig.getUntrackedParameter<edm::InputTag>("unsubjet_map", {});
if (!unsubjet_map_tag.label().empty()) {
unsubjet_map_token_ = consumes<JetMatchMap>(unsubjet_map_tag);
use_unsubjet_map_ = true;
}

produces<DeepBoostedJetTagInfoCollection>();
}

Expand All @@ -324,6 +335,7 @@ void DeepBoostedJetTagInfoProducer::fillDescriptions(edm::ConfigurationDescripti
desc.add<edm::InputTag>("secondary_vertices", edm::InputTag("inclusiveCandidateSecondaryVertices"));
desc.add<edm::InputTag>("pf_candidates", edm::InputTag("particleFlow"));
desc.add<edm::InputTag>("jets", edm::InputTag("ak8PFJetsPuppi"));
desc.addUntracked<edm::InputTag>("unsubjet_map", {});
desc.add<edm::InputTag>("puppi_value_map", edm::InputTag("puppi"));
desc.add<edm::InputTag>("vertex_associator", edm::InputTag("primaryVertexAssociation", "original"));
desc.add<bool>("use_scouting_features", false);
Expand All @@ -347,6 +359,7 @@ void DeepBoostedJetTagInfoProducer::produce(edm::Event &iEvent, const edm::Event
auto output_tag_infos = std::make_unique<DeepBoostedJetTagInfoCollection>();
// Input jets
auto jets = iEvent.getHandle(jet_token_);
auto unsubjet_map = use_unsubjet_map_ ? iEvent.getHandle(unsubjet_map_token_) : edm::Handle<JetMatchMap>();
// Primary vertexes
if (!use_scouting_features_) {
iEvent.getByToken(vtx_token_, vtxs_);
Expand Down Expand Up @@ -392,6 +405,8 @@ void DeepBoostedJetTagInfoProducer::produce(edm::Event &iEvent, const edm::Event
for (std::size_t jet_n = 0; jet_n < jets->size(); jet_n++) {
const auto &jet = (*jets)[jet_n];
edm::RefToBase<reco::Jet> jet_ref(jets, jet_n);
const auto &unsubJet =
(use_unsubjet_map_ && (*unsubjet_map)[jet_ref].isNonnull()) ? *(*unsubjet_map)[jet_ref] : jet;

// create jet features
DeepBoostedJetFeatures features;
Expand Down Expand Up @@ -421,13 +436,13 @@ void DeepBoostedJetTagInfoProducer::produce(edm::Event &iEvent, const edm::Event
if (jet.pt() < min_jet_pt_ or std::abs(jet.eta()) > max_jet_eta_) {
fill_vars = false;
}
if (jet.numberOfDaughters() == 0 and !use_scouting_features_) {
if (unsubJet.numberOfDaughters() == 0 and !use_scouting_features_) {
fill_vars = false;
}

// fill features
if (fill_vars) {
fillParticleFeatures(features, jet);
fillParticleFeatures(features, unsubJet, jet);
if (!use_scouting_features_) {
fillSVFeatures(features, jet);
}
Expand Down Expand Up @@ -475,7 +490,9 @@ bool DeepBoostedJetTagInfoProducer::useTrackProperties(const reco::PFCandidate *
return track != nullptr and track->pt() > min_pt_for_track_properties_;
};

void DeepBoostedJetTagInfoProducer::fillParticleFeatures(DeepBoostedJetFeatures &fts, const reco::Jet &jet) {
void DeepBoostedJetTagInfoProducer::fillParticleFeatures(DeepBoostedJetFeatures &fts,
const reco::Jet &unsubJet,
const reco::Jet &jet) {
// some jet properties
math::XYZVector jet_dir = jet.momentum().Unit();
TVector3 jet_direction(jet.momentum().Unit().x(), jet.momentum().Unit().y(), jet.momentum().Unit().z());
Expand All @@ -492,7 +509,7 @@ void DeepBoostedJetTagInfoProducer::fillParticleFeatures(DeepBoostedJetFeatures

// make list of pf-candidates to be considered
std::vector<reco::CandidatePtr> daughters;
for (const auto &dau : jet.daughterPtrVector()) {
for (const auto &dau : unsubJet.daughterPtrVector()) {
// remove particles w/ extremely low puppi weights
// [Note] use jet daughters here to get the puppiWgt correctly
if ((puppiWgt(dau)) < min_puppi_wgt_)
Expand Down
26 changes: 22 additions & 4 deletions RecoBTag/FeatureTools/plugins/DeepFlavourTagInfoProducer.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@ class HistogramProbabilityEstimator;
#include "FWCore/Framework/interface/EventSetupRecord.h"
#include "FWCore/Framework/interface/EventSetupRecordImplementation.h"
#include "FWCore/Framework/interface/EventSetupRecordKey.h"
#include "DataFormats/Common/interface/AssociationMap.h"

class DeepFlavourTagInfoProducer : public edm::stream::EDProducer<> {
public:
Expand All @@ -70,6 +71,7 @@ class DeepFlavourTagInfoProducer : public edm::stream::EDProducer<> {
typedef reco::VertexCompositePtrCandidateCollection SVCollection;
typedef reco::VertexCollection VertexCollection;
typedef edm::View<reco::ShallowTagInfo> ShallowTagInfoCollection;
typedef edm::AssociationMap<edm::OneToOne<reco::JetView, reco::JetView>> JetMatchMap;

void beginStream(edm::StreamID) override {}
void produce(edm::Event&, const edm::EventSetup&) override;
Expand All @@ -80,6 +82,7 @@ class DeepFlavourTagInfoProducer : public edm::stream::EDProducer<> {
const bool flip_;

edm::EDGetTokenT<edm::View<reco::Jet>> jet_token_;
edm::EDGetTokenT<JetMatchMap> unsubjet_map_token_;
edm::EDGetTokenT<VertexCollection> vtx_token_;
edm::EDGetTokenT<SVCollection> sv_token_;
edm::EDGetTokenT<ShallowTagInfoCollection> shallow_tag_info_token_;
Expand All @@ -93,6 +96,7 @@ class DeepFlavourTagInfoProducer : public edm::stream::EDProducer<> {

bool use_puppi_value_map_;
bool use_pvasq_value_map_;
bool use_unsubjet_map_;

bool fallback_puppi_weight_;
bool fallback_vertex_association_;
Expand Down Expand Up @@ -126,6 +130,7 @@ DeepFlavourTagInfoProducer::DeepFlavourTagInfoProducer(const edm::ParameterSet&
esConsumes<TransientTrackBuilder, TransientTrackRecord>(edm::ESInputTag("", "TransientTrackBuilder"))),
use_puppi_value_map_(false),
use_pvasq_value_map_(false),
use_unsubjet_map_(false),
fallback_puppi_weight_(iConfig.getParameter<bool>("fallback_puppi_weight")),
fallback_vertex_association_(iConfig.getParameter<bool>("fallback_vertex_association")),
run_deepVertex_(iConfig.getParameter<bool>("run_deepVertex")),
Expand Down Expand Up @@ -154,6 +159,12 @@ DeepFlavourTagInfoProducer::DeepFlavourTagInfoProducer(const edm::ParameterSet&
calib2d_token_ = esConsumes<TrackProbabilityCalibration, BTagTrackProbability2DRcd>();
calib3d_token_ = esConsumes<TrackProbabilityCalibration, BTagTrackProbability3DRcd>();
}

const auto& unsubjet_map_tag = iConfig.getUntrackedParameter<edm::InputTag>("unsubjet_map", {});
if (!unsubjet_map_tag.label().empty()) {
unsubjet_map_token_ = consumes<JetMatchMap>(unsubjet_map_tag);
use_unsubjet_map_ = true;
}
}

DeepFlavourTagInfoProducer::~DeepFlavourTagInfoProducer() {}
Expand All @@ -169,6 +180,7 @@ void DeepFlavourTagInfoProducer::fillDescriptions(edm::ConfigurationDescriptions
desc.add<edm::InputTag>("puppi_value_map", edm::InputTag("puppi"));
desc.add<edm::InputTag>("secondary_vertices", edm::InputTag("inclusiveCandidateSecondaryVertices"));
desc.add<edm::InputTag>("jets", edm::InputTag("ak4PFJetsCHS"));
desc.addUntracked<edm::InputTag>("unsubjet_map", {});
desc.add<edm::InputTag>("candidates", edm::InputTag("packedPFCandidates"));
desc.add<edm::InputTag>("vertex_associator", edm::InputTag("primaryVertexAssociation", "original"));
desc.add<bool>("fallback_puppi_weight", false);
Expand All @@ -189,6 +201,10 @@ void DeepFlavourTagInfoProducer::produce(edm::Event& iEvent, const edm::EventSet
edm::Handle<edm::View<reco::Jet>> jets;
iEvent.getByToken(jet_token_, jets);

edm::Handle<JetMatchMap> unsubjet_map;
if (use_unsubjet_map_)
iEvent.getByToken(unsubjet_map_token_, unsubjet_map);

edm::Handle<VertexCollection> vtxs;
iEvent.getByToken(vtx_token_, vtxs);
if (vtxs->empty()) {
Expand Down Expand Up @@ -256,6 +272,8 @@ void DeepFlavourTagInfoProducer::produce(edm::Event& iEvent, const edm::EventSet
const auto* pf_jet = dynamic_cast<const reco::PFJet*>(&jet);
const auto* pat_jet = dynamic_cast<const pat::Jet*>(&jet);
edm::RefToBase<reco::Jet> jet_ref(jets, jet_n);
const auto& unsubJet =
(use_unsubjet_map_ && (*unsubjet_map)[jet_ref].isNonnull()) ? *(*unsubjet_map)[jet_ref] : jet;
// TagInfoCollection not in an associative container so search for matchs
const edm::View<reco::ShallowTagInfo>& taginfos = *shallow_tag_infos;
edm::Ptr<reco::ShallowTagInfo> match;
Expand Down Expand Up @@ -316,8 +334,8 @@ void DeepFlavourTagInfoProducer::produce(edm::Event& iEvent, const edm::EventSet
// unsorted reference to sv
const auto& svs_unsorted = *svs;
// fill collection, from DeepTNtuples plus some styling
for (unsigned int i = 0; i < jet.numberOfDaughters(); i++) {
auto cand = jet.daughter(i);
for (unsigned int i = 0; i < unsubJet.numberOfDaughters(); i++) {
auto cand = unsubJet.daughter(i);
if (cand) {
// candidates under 950MeV (configurable) are not considered
// might change if we use also white-listing
Expand Down Expand Up @@ -350,9 +368,9 @@ void DeepFlavourTagInfoProducer::produce(edm::Event& iEvent, const edm::EventSet
features.n_pf_features.clear();
features.n_pf_features.resize(n_sorted.size());

for (unsigned int i = 0; i < jet.numberOfDaughters(); i++) {
for (unsigned int i = 0; i < unsubJet.numberOfDaughters(); i++) {
// get pointer and check that is correct
auto cand = dynamic_cast<const reco::Candidate*>(jet.daughter(i));
auto cand = dynamic_cast<const reco::Candidate*>(unsubJet.daughter(i));
if (!cand)
continue;
// candidates under 950MeV are not considered
Expand Down
26 changes: 22 additions & 4 deletions RecoBTag/FeatureTools/plugins/ParticleTransformerAK4TagInfoProducer.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -56,6 +56,7 @@ class HistogramProbabilityEstimator;
#include "FWCore/Framework/interface/EventSetupRecord.h"
#include "FWCore/Framework/interface/EventSetupRecordImplementation.h"
#include "FWCore/Framework/interface/EventSetupRecordKey.h"
#include "DataFormats/Common/interface/AssociationMap.h"

class ParticleTransformerAK4TagInfoProducer : public edm::stream::EDProducer<> {
public:
Expand All @@ -68,6 +69,7 @@ class ParticleTransformerAK4TagInfoProducer : public edm::stream::EDProducer<> {
typedef std::vector<reco::ParticleTransformerAK4TagInfo> ParticleTransformerAK4TagInfoCollection;
typedef reco::VertexCompositePtrCandidateCollection SVCollection;
typedef reco::VertexCollection VertexCollection;
typedef edm::AssociationMap<edm::OneToOne<reco::JetView, reco::JetView>> JetMatchMap;

void beginStream(edm::StreamID) override {}
void produce(edm::Event&, const edm::EventSetup&) override;
Expand All @@ -80,13 +82,15 @@ class ParticleTransformerAK4TagInfoProducer : public edm::stream::EDProducer<> {
const edm::EDGetTokenT<edm::View<reco::Jet>> jet_token_;
const edm::EDGetTokenT<VertexCollection> vtx_token_;
const edm::EDGetTokenT<SVCollection> sv_token_;
edm::EDGetTokenT<JetMatchMap> unsubjet_map_token_;
edm::EDGetTokenT<edm::ValueMap<float>> puppi_value_map_token_;
edm::EDGetTokenT<edm::ValueMap<int>> pvasq_value_map_token_;
edm::EDGetTokenT<edm::Association<VertexCollection>> pvas_token_;
const edm::EDGetTokenT<edm::View<reco::Candidate>> candidateToken_;
const edm::ESGetToken<TransientTrackBuilder, TransientTrackRecord> track_builder_token_;
bool use_puppi_value_map_;
bool use_pvasq_value_map_;
bool use_unsubjet_map_;

const bool fallback_puppi_weight_;
const bool fallback_vertex_association_;
Expand All @@ -109,6 +113,7 @@ ParticleTransformerAK4TagInfoProducer::ParticleTransformerAK4TagInfoProducer(con
esConsumes<TransientTrackBuilder, TransientTrackRecord>(edm::ESInputTag("", "TransientTrackBuilder"))),
use_puppi_value_map_(false),
use_pvasq_value_map_(false),
use_unsubjet_map_(false),
fallback_puppi_weight_(iConfig.getParameter<bool>("fallback_puppi_weight")),
fallback_vertex_association_(iConfig.getParameter<bool>("fallback_vertex_association")),
is_weighted_jet_(iConfig.getParameter<bool>("is_weighted_jet")),
Expand All @@ -131,6 +136,12 @@ ParticleTransformerAK4TagInfoProducer::ParticleTransformerAK4TagInfoProducer(con
pvas_token_ = consumes<edm::Association<VertexCollection>>(pvas_tag);
use_pvasq_value_map_ = true;
}

const auto& unsubjet_map_tag = iConfig.getUntrackedParameter<edm::InputTag>("unsubjet_map", {});
if (!unsubjet_map_tag.label().empty()) {
unsubjet_map_token_ = consumes<JetMatchMap>(unsubjet_map_tag);
use_unsubjet_map_ = true;
}
}

void ParticleTransformerAK4TagInfoProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
Expand All @@ -143,6 +154,7 @@ void ParticleTransformerAK4TagInfoProducer::fillDescriptions(edm::ConfigurationD
desc.add<edm::InputTag>("puppi_value_map", edm::InputTag("puppi"));
desc.add<edm::InputTag>("secondary_vertices", edm::InputTag("inclusiveCandidateSecondaryVertices"));
desc.add<edm::InputTag>("jets", edm::InputTag("ak4PFJetsCHS"));
desc.addUntracked<edm::InputTag>("unsubjet_map", {});
desc.add<edm::InputTag>("candidates", edm::InputTag("packedPFCandidates"));
desc.add<edm::InputTag>("vertex_associator", edm::InputTag("primaryVertexAssociation", "original"));
desc.add<bool>("fallback_puppi_weight", false);
Expand All @@ -158,6 +170,10 @@ void ParticleTransformerAK4TagInfoProducer::produce(edm::Event& iEvent, const ed
edm::Handle<edm::View<reco::Jet>> jets;
iEvent.getByToken(jet_token_, jets);

edm::Handle<JetMatchMap> unsubjet_map;
if (use_unsubjet_map_)
iEvent.getByToken(unsubjet_map_token_, unsubjet_map);

edm::Handle<VertexCollection> vtxs;
iEvent.getByToken(vtx_token_, vtxs);
if (vtxs->empty()) {
Expand Down Expand Up @@ -204,6 +220,8 @@ void ParticleTransformerAK4TagInfoProducer::produce(edm::Event& iEvent, const ed
const auto* pf_jet = dynamic_cast<const reco::PFJet*>(&jet);
const auto* pat_jet = dynamic_cast<const pat::Jet*>(&jet);
edm::RefToBase<reco::Jet> jet_ref(jets, jet_n);
const auto& unsubJet =
(use_unsubjet_map_ && (*unsubjet_map)[jet_ref].isNonnull()) ? *(*unsubjet_map)[jet_ref] : jet;

if (features.is_filled) {
math::XYZVector jet_dir = jet.momentum().Unit();
Expand Down Expand Up @@ -237,8 +255,8 @@ void ParticleTransformerAK4TagInfoProducer::produce(edm::Event& iEvent, const ed
// unsorted reference to sv
const auto& svs_unsorted = *svs;
// fill collection, from DeepTNtuples plus some styling
for (unsigned int i = 0; i < jet.numberOfDaughters(); i++) {
auto cand = jet.daughter(i);
for (unsigned int i = 0; i < unsubJet.numberOfDaughters(); i++) {
auto cand = unsubJet.daughter(i);
if (cand) {
// candidates under 950MeV (configurable) are not considered
// might change if we use also white-listing
Expand Down Expand Up @@ -274,9 +292,9 @@ void ParticleTransformerAK4TagInfoProducer::produce(edm::Event& iEvent, const ed
features.n_pf_features.clear();
features.n_pf_features.resize(n_sorted.size());

for (unsigned int i = 0; i < jet.numberOfDaughters(); i++) {
for (unsigned int i = 0; i < unsubJet.numberOfDaughters(); i++) {
// get pointer and check that is correct
auto cand = dynamic_cast<const reco::Candidate*>(jet.daughter(i));
auto cand = dynamic_cast<const reco::Candidate*>(unsubJet.daughter(i));
if (!cand)
continue;
// candidates under 950MeV are not considered
Expand Down
1 change: 1 addition & 0 deletions RecoBTag/FeatureTools/test/BuildFile.xml
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
<test name="testRecoBTagInfoWithUnsubJets" command="cmsRun ${LOCALTOP}/src/RecoBTag/FeatureTools/test/u0_cfg.py"/>
Loading

0 comments on commit 4e1f0e4

Please sign in to comment.