FeatureTools/plugins/DeepDoubleXTagInfoProducer.cc

0001
0002 #include "FWCore/Framework/interface/Frameworkfwd.h"
0003 #include "FWCore/Framework/interface/stream/EDProducer.h"
0004
0005 #include "FWCore/Framework/interface/Event.h"
0006 #include "FWCore/Framework/interface/MakerMacros.h"
0007
0008 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0009 #include "FWCore/Utilities/interface/StreamID.h"
0010 #include "FWCore/Utilities/interface/ESGetToken.h"
0011
0012 #include "DataFormats/Candidate/interface/Candidate.h"
0013
0014 #include "DataFormats/PatCandidates/interface/Jet.h"
0015 #include "DataFormats/PatCandidates/interface/PackedCandidate.h"
0016
0017 #include "DataFormats/BTauReco/interface/BoostedDoubleSVTagInfo.h"
0018
0019 #include "TrackingTools/Records/interface/TransientTrackRecord.h"
0020 #include "TrackingTools/TransientTrack/interface/TransientTrackBuilder.h"
0021
0022 #include "DataFormats/BTauReco/interface/DeepDoubleXFeatures.h"
0023 #include "DataFormats/BTauReco/interface/DeepDoubleXTagInfo.h"
0024
0025 #include "RecoBTag/FeatureTools/interface/BoostedDoubleSVTagInfoConverter.h"
0026 #include "RecoBTag/FeatureTools/interface/NeutralCandidateConverter.h"
0027 #include "RecoBTag/FeatureTools/interface/ChargedCandidateConverter.h"
0028 #include "RecoBTag/FeatureTools/interface/JetConverter.h"
0029 #include "RecoBTag/FeatureTools/interface/SecondaryVertexConverter.h"
0030
0031 #include "RecoBTag/FeatureTools/interface/TrackInfoBuilder.h"
0032 #include "RecoBTag/FeatureTools/interface/sorting_modules.h"
0033
0034 #include "DataFormats/Candidate/interface/VertexCompositePtrCandidate.h"
0035 #include "DataFormats/VertexReco/interface/VertexFwd.h"
0036
0037 #include "RecoBTag/FeatureTools/interface/deep_helpers.h"
0038
0039 class DeepDoubleXTagInfoProducer : public edm::stream::EDProducer<> {
0040 public:
0041   explicit DeepDoubleXTagInfoProducer(const edm::ParameterSet&);
0042   ~DeepDoubleXTagInfoProducer() override;
0043
0044   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0045
0046 private:
0047   typedef std::vector<reco::DeepDoubleXTagInfo> DeepDoubleXTagInfoCollection;
0048   typedef reco::VertexCompositePtrCandidateCollection SVCollection;
0049   typedef reco::VertexCollection VertexCollection;
0050   typedef edm::View<reco::BoostedDoubleSVTagInfo> BoostedDoubleSVTagInfoCollection;
0051
0052   void beginStream(edm::StreamID) override {}
0053   void produce(edm::Event&, const edm::EventSetup&) override;
0054   void endStream() override {}
0055
0056   const double jet_radius_;
0057   const double min_jet_pt_;
0058   const double min_candidate_pt_;
0059
0060   edm::EDGetTokenT<edm::View<reco::Jet>> jet_token_;
0061   edm::EDGetTokenT<VertexCollection> vtx_token_;
0062   edm::EDGetTokenT<SVCollection> sv_token_;
0063   edm::EDGetTokenT<BoostedDoubleSVTagInfoCollection> shallow_tag_info_token_;
0064   edm::ESGetToken<TransientTrackBuilder, TransientTrackRecord> track_builder_token_;
0065 };
0066
0067 DeepDoubleXTagInfoProducer::DeepDoubleXTagInfoProducer(const edm::ParameterSet& iConfig)
0068     : jet_radius_(iConfig.getParameter<double>("jet_radius")),
0069       min_jet_pt_(iConfig.getParameter<double>("min_jet_pt")),
0070       min_candidate_pt_(iConfig.getParameter<double>("min_candidate_pt")),
0071       jet_token_(consumes<edm::View<reco::Jet>>(iConfig.getParameter<edm::InputTag>("jets"))),
0072       vtx_token_(consumes<VertexCollection>(iConfig.getParameter<edm::InputTag>("vertices"))),
0073       sv_token_(consumes<SVCollection>(iConfig.getParameter<edm::InputTag>("secondary_vertices"))),
0074       shallow_tag_info_token_(
0075           consumes<BoostedDoubleSVTagInfoCollection>(iConfig.getParameter<edm::InputTag>("shallow_tag_infos"))),
0076       track_builder_token_(
0077           esConsumes<TransientTrackBuilder, TransientTrackRecord>(edm::ESInputTag("", "TransientTrackBuilder"))) {
0078   produces<DeepDoubleXTagInfoCollection>();
0079 }
0080
0081 DeepDoubleXTagInfoProducer::~DeepDoubleXTagInfoProducer() {}
0082
0083 void DeepDoubleXTagInfoProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0084   // pfDeepDoubleXTagInfos
0085   edm::ParameterSetDescription desc;
0086   desc.add<edm::InputTag>("shallow_tag_infos", edm::InputTag("pfBoostedDoubleSVAK8TagInfos"));
0087   desc.add<double>("jet_radius", 0.8);
0088   desc.add<double>("min_jet_pt", 150);
0089   desc.add<double>("min_candidate_pt", 0.95);
0090   desc.add<edm::InputTag>("vertices", edm::InputTag("offlinePrimaryVertices"));
0091   desc.add<edm::InputTag>("secondary_vertices", edm::InputTag("inclusiveCandidateSecondaryVertices"));
0092   desc.add<edm::InputTag>("jets", edm::InputTag("ak8PFJetsPuppi"));
0093   descriptions.add("pfDeepDoubleXTagInfos", desc);
0094 }
0095
0096 void DeepDoubleXTagInfoProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
0097   auto output_tag_infos = std::make_unique<DeepDoubleXTagInfoCollection>();
0098
0099   edm::Handle<edm::View<reco::Jet>> jets;
0100   iEvent.getByToken(jet_token_, jets);
0101
0102   edm::Handle<VertexCollection> vtxs;
0103   iEvent.getByToken(vtx_token_, vtxs);
0104   if (vtxs->empty()) {
0105     // produce empty TagInfos in case no primary vertex
0106     iEvent.put(std::move(output_tag_infos));
0107     return;  // exit event
0108   }
0109   // reference to primary vertex
0110   const auto& pv = vtxs->at(0);
0111
0112   edm::Handle<SVCollection> svs;
0113   iEvent.getByToken(sv_token_, svs);
0114
0115   edm::Handle<BoostedDoubleSVTagInfoCollection> shallow_tag_infos;
0116   iEvent.getByToken(shallow_tag_info_token_, shallow_tag_infos);
0117
0118   edm::ESHandle<TransientTrackBuilder> track_builder = iSetup.getHandle(track_builder_token_);
0119
0120   for (std::size_t jet_n = 0; jet_n < jets->size(); jet_n++) {
0121     // create data containing structure
0122     btagbtvdeep::DeepDoubleXFeatures features;
0123
0124     // reco jet reference (use as much as possible)
0125     const auto& jet = jets->at(jet_n);
0126     const auto* pf_jet = dynamic_cast<const reco::PFJet*>(&jet);
0127     const auto* pat_jet = dynamic_cast<const pat::Jet*>(&jet);
0128
0129     edm::RefToBase<reco::Jet> jet_ref(jets, jet_n);
0130     if (jet.pt() > min_jet_pt_) {
0131       features.filled();
0132       // TagInfoCollection not in an associative container so search for matchs
0133       const edm::View<reco::BoostedDoubleSVTagInfo>& taginfos = *shallow_tag_infos;
0134       edm::Ptr<reco::BoostedDoubleSVTagInfo> match;
0135       // Try first by 'same index'
0136       if ((jet_n < taginfos.size()) && (taginfos[jet_n].jet() == jet_ref)) {
0137         match = taginfos.ptrAt(jet_n);
0138       } else {
0139         // otherwise fall back to a simple search
0140         for (auto itTI = taginfos.begin(), edTI = taginfos.end(); itTI != edTI; ++itTI) {
0141           if (itTI->jet() == jet_ref) {
0142             match = taginfos.ptrAt(itTI - taginfos.begin());
0143             break;
0144           }
0145         }
0146       }
0147       reco::BoostedDoubleSVTagInfo tag_info;
0148       if (match.isNonnull()) {
0149         tag_info = *match;
0150       }  // will be default values otherwise
0151
0152       // fill basic jet features
0153       btagbtvdeep::JetConverter::jetToFeatures(jet, features.jet_features);
0154
0155       // fill number of pv
0156       features.npv = vtxs->size();
0157
0158       // fill features from BoostedDoubleSVTagInfo
0159       const auto& tag_info_vars = tag_info.taggingVariables();
0160       btagbtvdeep::doubleBTagToFeatures(tag_info_vars, features.tag_info_features);
0161
0162       // copy which will be sorted
0163       auto svs_sorted = *svs;
0164       // sort by dxy
0165       std::sort(svs_sorted.begin(), svs_sorted.end(), [&pv](const auto& sva, const auto& svb) {
0166         return btagbtvdeep::sv_vertex_comparator(sva, svb, pv);
0167       });
0168       // fill features from secondary vertices
0169       for (const auto& sv : svs_sorted) {
0170         if (reco::deltaR(sv, jet) > jet_radius_)
0171           continue;
0172         else {
0173           features.sv_features.emplace_back();
0174           // in C++17 could just get from emplace_back output
0175           auto& sv_features = features.sv_features.back();
0176           btagbtvdeep::svToFeatures(sv, pv, jet, sv_features);
0177         }
0178       }
0179
0180       // stuff required for dealing with pf candidates
0181       math::XYZVector jet_dir = jet.momentum().Unit();
0182       GlobalVector jet_ref_track_dir(jet.px(), jet.py(), jet.pz());
0183
0184       std::vector<btagbtvdeep::SortingClass<size_t>> c_sorted, n_sorted;
0185       std::vector<int> n_indexes;
0186
0187       // to cache the TrackInfo
0188       std::map<unsigned int, btagbtvdeep::TrackInfoBuilder> trackinfos;
0189
0190       // unsorted reference to sv
0191       const auto& svs_unsorted = *svs;
0192       // fill collection, from DeepTNtuples plus some styling
0193       // std::vector<const pat::PackedCandidate*> daughters;
0194       std::vector<const reco::Candidate*> daughters;
0195       for (unsigned int i = 0; i < jet.numberOfDaughters(); i++) {
0196         auto const* cand = jet.daughter(i);
0197         auto packed_cand = dynamic_cast<const pat::PackedCandidate*>(cand);
0198         auto reco_cand = dynamic_cast<const reco::PFCandidate*>(cand);
0199         if (packed_cand) {
0200           if (cand->numberOfDaughters() > 0) {
0201             for (unsigned int k = 0; k < cand->numberOfDaughters(); k++) {
0202               daughters.push_back(dynamic_cast<const pat::PackedCandidate*>(cand->daughter(k)));
0203             }
0204           } else {
0205             auto packed_cand = dynamic_cast<const pat::PackedCandidate*>(cand);
0206             daughters.push_back(packed_cand);
0207           }
0208         } else if (reco_cand) {
0209           // need some edm::Ptr or edm::Ref if reco candidates
0210           // dynamical casting to pointers, null if not possible
0211           daughters.push_back(reco_cand);
0212         }
0213       }
0214
0215       std::sort(daughters.begin(), daughters.end(), [](const auto& a, const auto& b) { return a->pt() > b->pt(); });
0216       for (unsigned int i = 0; i < daughters.size(); i++) {
0217         auto const* cand = daughters.at(i);
0218         if (cand) {
0219           // candidates under 950MeV (configurable) are not considered
0220           // might change if we use also white-listing
0221           if (cand->pt() < min_candidate_pt_)
0222             continue;
0223           if (cand->charge() != 0) {
0224             auto& trackinfo = trackinfos.emplace(i, track_builder).first->second;
0225             trackinfo.buildTrackInfo(cand, jet_dir, jet_ref_track_dir, pv);
0226             c_sorted.emplace_back(i,
0227                                   trackinfo.getTrackSip2dSig(),
0228                                   -btagbtvdeep::mindrsvpfcand(svs_unsorted, cand, jet_radius_),
0229                                   cand->pt() / jet.pt());
0230           } else {
0231             n_sorted.emplace_back(
0232                 i, -1, -btagbtvdeep::mindrsvpfcand(svs_unsorted, cand, jet_radius_), cand->pt() / jet.pt());
0233             n_indexes.push_back(i);
0234           }
0235         }
0236       }
0237
0238       // sort collections in added order of priority
0239       std::sort(c_sorted.begin(), c_sorted.end(), btagbtvdeep::SortingClass<std::size_t>::compareByABCInv);
0240       std::sort(n_sorted.begin(), n_sorted.end(), btagbtvdeep::SortingClass<std::size_t>::compareByABCInv);
0241
0242       std::vector<size_t> c_sortedindices, n_sortedindices;
0243
0244       // this puts 0 everywhere and the right position in ind
0245       c_sortedindices = btagbtvdeep::invertSortingVector(c_sorted);
0246       n_sortedindices = btagbtvdeep::invertSortingVector(n_sorted);
0247
0248       // set right size to vectors
0249       features.c_pf_features.clear();
0250       features.c_pf_features.resize(c_sorted.size());
0251       features.n_pf_features.clear();
0252       features.n_pf_features.resize(n_sorted.size());
0253
0254       for (unsigned int i = 0; i < daughters.size(); i++) {
0255         auto const* cand = daughters.at(i);
0256         if (cand) {
0257           // candidates under 950MeV are not considered
0258           // might change if we use also white-listing
0259           if (cand->pt() < min_candidate_pt_)
0260             continue;
0261           auto packed_cand = dynamic_cast<const pat::PackedCandidate*>(cand);
0262           auto reco_cand = dynamic_cast<const reco::PFCandidate*>(cand);
0263
0264           // need some edm::Ptr or edm::Ref if reco candidates
0265           reco::PFCandidatePtr reco_ptr;
0266           if (pf_jet) {
0267             reco_ptr = pf_jet->getPFConstituent(i);
0268           } else if (pat_jet && reco_cand) {
0269             reco_ptr = pat_jet->getPFConstituent(i);
0270           }
0271
0272           float puppiw = 1.0;  // fallback value
0273
0274           float drminpfcandsv = btagbtvdeep::mindrsvpfcand(svs_unsorted, cand, jet_radius_);
0275           if (cand->charge() != 0) {
0276             // is charged candidate
0277             auto entry = c_sortedindices.at(i);
0278             // get cached track info
0279             auto& trackinfo = trackinfos.at(i);
0280             // get_ref to vector element
0281             auto& c_pf_features = features.c_pf_features.at(entry);
0282             if (packed_cand) {
0283               btagbtvdeep::packedCandidateToFeatures(
0284                   packed_cand, *pat_jet, trackinfo, drminpfcandsv, static_cast<float>(jet_radius_), c_pf_features);
0285             } else if (reco_cand) {
0286               // get vertex association quality
0287               int pv_ass_quality = 0;  // fallback value
0288               // getting the PV as PackedCandidatesProducer
0289               // but using not the slimmed but original vertices
0290               auto ctrack = reco_cand->bestTrack();
0291               int pvi = -1;
0292               float dist = 1e99;
0293               for (size_t ii = 0; ii < vtxs->size(); ii++) {
0294                 float dz = (ctrack) ? std::abs(ctrack->dz(((*vtxs)[ii]).position())) : 0;
0295                 if (dz < dist) {
0296                   pvi = ii;
0297                   dist = dz;
0298                 }
0299               }
0300               auto pv = reco::VertexRef(vtxs, pvi);
0301               btagbtvdeep::recoCandidateToFeatures(reco_cand,
0302                                                    jet,
0303                                                    trackinfo,
0304                                                    drminpfcandsv,
0305                                                    static_cast<float>(jet_radius_),
0306                                                    puppiw,
0307                                                    pv_ass_quality,
0308                                                    pv,
0309                                                    c_pf_features);
0310             }
0311           } else {
0312             // is neutral candidate
0313             auto entry = n_sortedindices.at(i);
0314             // get_ref to vector element
0315             auto& n_pf_features = features.n_pf_features.at(entry);
0316             // // fill feature structure
0317             if (packed_cand) {
0318               btagbtvdeep::packedCandidateToFeatures(
0319                   packed_cand, *pat_jet, drminpfcandsv, static_cast<float>(jet_radius_), n_pf_features);
0320             } else if (reco_cand) {
0321               btagbtvdeep::recoCandidateToFeatures(
0322                   reco_cand, jet, drminpfcandsv, static_cast<float>(jet_radius_), puppiw, n_pf_features);
0323             }
0324           }
0325         }
0326       }
0327     }
0328     output_tag_infos->emplace_back(features, jet_ref);
0329   }
0330
0331   iEvent.put(std::move(output_tag_infos));
0332 }
0333
0334 // define this as a plug-in
0335 DEFINE_FWK_MODULE(DeepDoubleXTagInfoProducer);