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* pat_jet = dynamic_cast<const pat::Jet*>(&jet);
0127     if (!pat_jet)
0128       throw edm::Exception(edm::errors::InvalidReference) << "Input is not a pat::Jet.";
0129
0130     edm::RefToBase<reco::Jet> jet_ref(jets, jet_n);
0131     if (jet.pt() > min_jet_pt_) {
0132       features.filled();
0133       // TagInfoCollection not in an associative container so search for matchs
0134       const edm::View<reco::BoostedDoubleSVTagInfo>& taginfos = *shallow_tag_infos;
0135       edm::Ptr<reco::BoostedDoubleSVTagInfo> match;
0136       // Try first by 'same index'
0137       if ((jet_n < taginfos.size()) && (taginfos[jet_n].jet() == jet_ref)) {
0138         match = taginfos.ptrAt(jet_n);
0139       } else {
0140         // otherwise fall back to a simple search
0141         for (auto itTI = taginfos.begin(), edTI = taginfos.end(); itTI != edTI; ++itTI) {
0142           if (itTI->jet() == jet_ref) {
0143             match = taginfos.ptrAt(itTI - taginfos.begin());
0144             break;
0145           }
0146         }
0147       }
0148       reco::BoostedDoubleSVTagInfo tag_info;
0149       if (match.isNonnull()) {
0150         tag_info = *match;
0151       }  // will be default values otherwise
0152
0153       // fill basic jet features
0154       btagbtvdeep::JetConverter::jetToFeatures(jet, features.jet_features);
0155
0156       // fill number of pv
0157       features.npv = vtxs->size();
0158
0159       // fill features from BoostedDoubleSVTagInfo
0160       const auto& tag_info_vars = tag_info.taggingVariables();
0161       btagbtvdeep::doubleBTagToFeatures(tag_info_vars, features.tag_info_features);
0162
0163       // copy which will be sorted
0164       auto svs_sorted = *svs;
0165       // sort by dxy
0166       std::sort(svs_sorted.begin(), svs_sorted.end(), [&pv](const auto& sva, const auto& svb) {
0167         return btagbtvdeep::sv_vertex_comparator(sva, svb, pv);
0168       });
0169       // fill features from secondary vertices
0170       for (const auto& sv : svs_sorted) {
0171         if (reco::deltaR(sv, jet) > jet_radius_)
0172           continue;
0173         else {
0174           features.sv_features.emplace_back();
0175           // in C++17 could just get from emplace_back output
0176           auto& sv_features = features.sv_features.back();
0177           btagbtvdeep::svToFeatures(sv, pv, jet, sv_features);
0178         }
0179       }
0180
0181       // stuff required for dealing with pf candidates
0182       math::XYZVector jet_dir = jet.momentum().Unit();
0183       GlobalVector jet_ref_track_dir(jet.px(), jet.py(), jet.pz());
0184
0185       std::vector<btagbtvdeep::SortingClass<size_t>> c_sorted, n_sorted;
0186       std::vector<int> n_indexes;
0187
0188       // to cache the TrackInfo
0189       std::map<unsigned int, btagbtvdeep::TrackInfoBuilder> trackinfos;
0190
0191       // unsorted reference to sv
0192       const auto& svs_unsorted = *svs;
0193       // fill collection, from DeepTNtuples plus some styling
0194       // std::vector<const pat::PackedCandidate*> daughters;
0195       std::vector<const reco::Candidate*> daughters;
0196       for (unsigned int i = 0; i < jet.numberOfDaughters(); i++) {
0197         auto const* cand = jet.daughter(i);
0198         auto packed_cand = dynamic_cast<const pat::PackedCandidate*>(cand);
0199         auto reco_cand = dynamic_cast<const reco::PFCandidate*>(cand);
0200         if (packed_cand) {
0201           if (cand->numberOfDaughters() > 0) {
0202             for (unsigned int k = 0; k < cand->numberOfDaughters(); k++) {
0203               daughters.push_back(dynamic_cast<const pat::PackedCandidate*>(cand->daughter(k)));
0204             }
0205           } else {
0206             auto packed_cand = dynamic_cast<const pat::PackedCandidate*>(cand);
0207             daughters.push_back(packed_cand);
0208           }
0209         } else if (reco_cand) {
0210           // need some edm::Ptr or edm::Ref if reco candidates
0211           // dynamical casting to pointers, null if not possible
0212           daughters.push_back(reco_cand);
0213         }
0214       }
0215
0216       std::sort(daughters.begin(), daughters.end(), [](const auto& a, const auto& b) { return a->pt() > b->pt(); });
0217       for (unsigned int i = 0; i < daughters.size(); i++) {
0218         auto const* cand = daughters.at(i);
0219         if (cand) {
0220           // candidates under 950MeV (configurable) are not considered
0221           // might change if we use also white-listing
0222           if (cand->pt() < min_candidate_pt_)
0223             continue;
0224           if (cand->charge() != 0) {
0225             auto& trackinfo = trackinfos.emplace(i, track_builder).first->second;
0226             trackinfo.buildTrackInfo(cand, jet_dir, jet_ref_track_dir, pv);
0227             c_sorted.emplace_back(i,
0228                                   trackinfo.getTrackSip2dSig(),
0229                                   -btagbtvdeep::mindrsvpfcand(svs_unsorted, cand, jet_radius_),
0230                                   cand->pt() / jet.pt());
0231           } else {
0232             n_sorted.emplace_back(
0233                 i, -1, -btagbtvdeep::mindrsvpfcand(svs_unsorted, cand, jet_radius_), cand->pt() / jet.pt());
0234             n_indexes.push_back(i);
0235           }
0236         }
0237       }
0238
0239       // sort collections in added order of priority
0240       std::sort(c_sorted.begin(), c_sorted.end(), btagbtvdeep::SortingClass<std::size_t>::compareByABCInv);
0241       std::sort(n_sorted.begin(), n_sorted.end(), btagbtvdeep::SortingClass<std::size_t>::compareByABCInv);
0242
0243       std::vector<size_t> c_sortedindices, n_sortedindices;
0244
0245       // this puts 0 everywhere and the right position in ind
0246       c_sortedindices = btagbtvdeep::invertSortingVector(c_sorted);
0247       n_sortedindices = btagbtvdeep::invertSortingVector(n_sorted);
0248
0249       // set right size to vectors
0250       features.c_pf_features.clear();
0251       features.c_pf_features.resize(c_sorted.size());
0252       features.n_pf_features.clear();
0253       features.n_pf_features.resize(n_sorted.size());
0254
0255       for (unsigned int i = 0; i < daughters.size(); i++) {
0256         auto const* cand = daughters.at(i);
0257         if (cand) {
0258           // candidates under 950MeV are not considered
0259           // might change if we use also white-listing
0260           if (cand->pt() < min_candidate_pt_)
0261             continue;
0262           auto packed_cand = dynamic_cast<const pat::PackedCandidate*>(cand);
0263           auto reco_cand = dynamic_cast<const reco::PFCandidate*>(cand);
0264
0265           float puppiw = 1.0;  // fallback value
0266
0267           float drminpfcandsv = btagbtvdeep::mindrsvpfcand(svs_unsorted, cand, jet_radius_);
0268           if (cand->charge() != 0) {
0269             // is charged candidate
0270             auto entry = c_sortedindices.at(i);
0271             // get cached track info
0272             auto& trackinfo = trackinfos.at(i);
0273             // get_ref to vector element
0274             auto& c_pf_features = features.c_pf_features.at(entry);
0275             if (packed_cand) {
0276               btagbtvdeep::packedCandidateToFeatures(
0277                   packed_cand, *pat_jet, trackinfo, drminpfcandsv, static_cast<float>(jet_radius_), c_pf_features);
0278             } else if (reco_cand) {
0279               // get vertex association quality
0280               int pv_ass_quality = 0;  // fallback value
0281               // getting the PV as PackedCandidatesProducer
0282               // but using not the slimmed but original vertices
0283               auto ctrack = reco_cand->bestTrack();
0284               int pvi = -1;
0285               float dist = 1e99;
0286               for (size_t ii = 0; ii < vtxs->size(); ii++) {
0287                 float dz = (ctrack) ? std::abs(ctrack->dz(((*vtxs)[ii]).position())) : 0;
0288                 if (dz < dist) {
0289                   pvi = ii;
0290                   dist = dz;
0291                 }
0292               }
0293               auto pv = reco::VertexRef(vtxs, pvi);
0294               btagbtvdeep::recoCandidateToFeatures(reco_cand,
0295                                                    jet,
0296                                                    trackinfo,
0297                                                    drminpfcandsv,
0298                                                    static_cast<float>(jet_radius_),
0299                                                    puppiw,
0300                                                    pv_ass_quality,
0301                                                    pv,
0302                                                    c_pf_features);
0303             }
0304           } else {
0305             // is neutral candidate
0306             auto entry = n_sortedindices.at(i);
0307             // get_ref to vector element
0308             auto& n_pf_features = features.n_pf_features.at(entry);
0309             // // fill feature structure
0310             if (packed_cand) {
0311               btagbtvdeep::packedCandidateToFeatures(
0312                   packed_cand, *pat_jet, drminpfcandsv, static_cast<float>(jet_radius_), n_pf_features);
0313             } else if (reco_cand) {
0314               btagbtvdeep::recoCandidateToFeatures(
0315                   reco_cand, jet, drminpfcandsv, static_cast<float>(jet_radius_), puppiw, n_pf_features);
0316             }
0317           }
0318         }
0319       }
0320     }
0321     output_tag_infos->emplace_back(features, jet_ref);
0322   }
0323
0324   iEvent.put(std::move(output_tag_infos));
0325 }
0326
0327 // define this as a plug-in
0328 DEFINE_FWK_MODULE(DeepDoubleXTagInfoProducer);