Merge pull request #44474 from stahlleiton/UnsubJetsForBTag_CMSSW_14_0_X

[14_0_X] Add use of unsubtracted jets in RecoBTag info producers

RecoBTag/FeatureTools/plugins/DeepBoostedJetTagInfoProducer.cc

@@ -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;
 
@@ -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);
@@ -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_;
@@ -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")),
@@ -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>();
 }
 
@@ -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);
@@ -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_);
@@ -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;
@@ -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);
       }
@@ -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());
@@ -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_)

RecoBTag/FeatureTools/plugins/DeepFlavourTagInfoProducer.cc

@@ -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:
@@ -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;
@@ -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_;
@@ -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_;
@@ -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")),
@@ -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() {}
@@ -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);
@@ -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()) {
@@ -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;
@@ -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
@@ -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

RecoBTag/FeatureTools/plugins/ParticleTransformerAK4TagInfoProducer.cc

@@ -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:
@@ -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;
@@ -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_;
@@ -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")),
@@ -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) {
@@ -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);
@@ -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()) {
@@ -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();
@@ -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
@@ -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

RecoBTag/FeatureTools/test/BuildFile.xml

@@ -0,0 +1 @@
+<test name="testRecoBTagInfoWithUnsubJets" command="cmsRun ${LOCALTOP}/src/RecoBTag/FeatureTools/test/u0_cfg.py"/>