FeatureTools/plugins/ParticleTransformerAK4TagInfoProducer.cc

0001 #include "FWCore/Framework/interface/Frameworkfwd.h"
0002 #include "FWCore/Framework/interface/stream/EDProducer.h"
0003
0004 #include "FWCore/Framework/interface/Event.h"
0005 #include "FWCore/Framework/interface/MakerMacros.h"
0006
0007 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0008 #include "FWCore/Utilities/interface/StreamID.h"
0009 #include "FWCore/Utilities/interface/ESGetToken.h"
0010
0011 #include "DataFormats/PatCandidates/interface/Jet.h"
0012 #include "DataFormats/PatCandidates/interface/PackedCandidate.h"
0013
0014 #include "DataFormats/BTauReco/interface/ShallowTagInfo.h"
0015
0016 #include "TrackingTools/TransientTrack/interface/TransientTrackBuilder.h"
0017 #include "TrackingTools/Records/interface/TransientTrackRecord.h"
0018
0019 #include "DataFormats/BTauReco/interface/ParticleTransformerAK4TagInfo.h"
0020 #include "DataFormats/BTauReco/interface/ParticleTransformerAK4Features.h"
0021
0022 #include "RecoBTag/FeatureTools/interface/JetConverter.h"
0023 #include "RecoBTag/FeatureTools/interface/ShallowTagInfoConverter.h"
0024 #include "RecoBTag/FeatureTools/interface/SecondaryVertexConverter.h"
0025 #include "RecoBTag/FeatureTools/interface/NeutralCandidateConverter.h"
0026 #include "RecoBTag/FeatureTools/interface/ChargedCandidateConverter.h"
0027
0028 #include "RecoBTag/FeatureTools/interface/TrackInfoBuilder.h"
0029 #include "RecoBTag/FeatureTools/interface/sorting_modules.h"
0030
0031 #include "DataFormats/VertexReco/interface/VertexFwd.h"
0032 #include "DataFormats/Candidate/interface/VertexCompositePtrCandidate.h"
0033
0034 #include "RecoBTag/FeatureTools/interface/deep_helpers.h"
0035
0036 #include "FWCore/ParameterSet/interface/Registry.h"
0037 #include "FWCore/Common/interface/Provenance.h"
0038 #include "DataFormats/Provenance/interface/ProductProvenance.h"
0039
0040 #include "DataFormats/BTauReco/interface/SeedingTrackFeatures.h"
0041 #include "DataFormats/BTauReco/interface/TrackPairFeatures.h"
0042 #include "RecoBTag/FeatureTools/interface/TrackPairInfoBuilder.h"
0043 #include "RecoBTag/FeatureTools/interface/SeedingTrackInfoBuilder.h"
0044 #include "RecoBTag/FeatureTools/interface/SeedingTracksConverter.h"
0045
0046 #include "FWCore/Framework/interface/ESHandle.h"
0047 #include "FWCore/Framework/interface/EventSetup.h"
0048 #include "RecoBTag/TrackProbability/interface/HistogramProbabilityEstimator.h"
0049 class HistogramProbabilityEstimator;
0050 #include <memory>
0051
0052 #include <typeinfo>
0053 #include "CondFormats/BTauObjects/interface/TrackProbabilityCalibration.h"
0054 #include "CondFormats/DataRecord/interface/BTagTrackProbability2DRcd.h"
0055 #include "CondFormats/DataRecord/interface/BTagTrackProbability3DRcd.h"
0056 #include "FWCore/Framework/interface/EventSetupRecord.h"
0057 #include "FWCore/Framework/interface/EventSetupRecordImplementation.h"
0058 #include "FWCore/Framework/interface/EventSetupRecordKey.h"
0059
0060 class ParticleTransformerAK4TagInfoProducer : public edm::stream::EDProducer<> {
0061 public:
0062   explicit ParticleTransformerAK4TagInfoProducer(const edm::ParameterSet&);
0063   ~ParticleTransformerAK4TagInfoProducer() override = default;
0064
0065   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0066
0067 private:
0068   typedef std::vector<reco::ParticleTransformerAK4TagInfo> ParticleTransformerAK4TagInfoCollection;
0069   typedef reco::VertexCompositePtrCandidateCollection SVCollection;
0070   typedef reco::VertexCollection VertexCollection;
0071
0072   void beginStream(edm::StreamID) override {}
0073   void produce(edm::Event&, const edm::EventSetup&) override;
0074   void endStream() override {}
0075
0076   const double jet_radius_;
0077   const double min_candidate_pt_;
0078   const bool flip_;
0079
0080   const edm::EDGetTokenT<edm::View<reco::Jet>> jet_token_;
0081   const edm::EDGetTokenT<VertexCollection> vtx_token_;
0082   const edm::EDGetTokenT<SVCollection> sv_token_;
0083   edm::EDGetTokenT<edm::ValueMap<float>> puppi_value_map_token_;
0084   edm::EDGetTokenT<edm::ValueMap<int>> pvasq_value_map_token_;
0085   edm::EDGetTokenT<edm::Association<VertexCollection>> pvas_token_;
0086   const edm::EDGetTokenT<edm::View<reco::Candidate>> candidateToken_;
0087   const edm::ESGetToken<TransientTrackBuilder, TransientTrackRecord> track_builder_token_;
0088   bool use_puppi_value_map_;
0089   bool use_pvasq_value_map_;
0090
0091   const bool fallback_puppi_weight_;
0092   const bool fallback_vertex_association_;
0093
0094   const bool is_weighted_jet_;
0095
0096   const double min_jet_pt_;
0097   const double max_jet_eta_;
0098 };
0099
0100 ParticleTransformerAK4TagInfoProducer::ParticleTransformerAK4TagInfoProducer(const edm::ParameterSet& iConfig)
0101     : jet_radius_(iConfig.getParameter<double>("jet_radius")),
0102       min_candidate_pt_(iConfig.getParameter<double>("min_candidate_pt")),
0103       flip_(iConfig.getParameter<bool>("flip")),
0104       jet_token_(consumes<edm::View<reco::Jet>>(iConfig.getParameter<edm::InputTag>("jets"))),
0105       vtx_token_(consumes<VertexCollection>(iConfig.getParameter<edm::InputTag>("vertices"))),
0106       sv_token_(consumes<SVCollection>(iConfig.getParameter<edm::InputTag>("secondary_vertices"))),
0107       candidateToken_(consumes<edm::View<reco::Candidate>>(iConfig.getParameter<edm::InputTag>("candidates"))),
0108       track_builder_token_(
0109           esConsumes<TransientTrackBuilder, TransientTrackRecord>(edm::ESInputTag("", "TransientTrackBuilder"))),
0110       use_puppi_value_map_(false),
0111       use_pvasq_value_map_(false),
0112       fallback_puppi_weight_(iConfig.getParameter<bool>("fallback_puppi_weight")),
0113       fallback_vertex_association_(iConfig.getParameter<bool>("fallback_vertex_association")),
0114       is_weighted_jet_(iConfig.getParameter<bool>("is_weighted_jet")),
0115       min_jet_pt_(iConfig.getParameter<double>("min_jet_pt")),
0116       max_jet_eta_(iConfig.getParameter<double>("max_jet_eta")) {
0117   produces<ParticleTransformerAK4TagInfoCollection>();
0118
0119   const auto& puppi_value_map_tag = iConfig.getParameter<edm::InputTag>("puppi_value_map");
0120   if (!puppi_value_map_tag.label().empty()) {
0121     puppi_value_map_token_ = consumes<edm::ValueMap<float>>(puppi_value_map_tag);
0122     use_puppi_value_map_ = true;
0123   } else if (is_weighted_jet_) {
0124     throw edm::Exception(edm::errors::Configuration,
0125                          "puppi_value_map is not set but jet is weighted. Must set puppi_value_map.");
0126   }
0127
0128   const auto& pvas_tag = iConfig.getParameter<edm::InputTag>("vertex_associator");
0129   if (!pvas_tag.label().empty()) {
0130     pvasq_value_map_token_ = consumes<edm::ValueMap<int>>(pvas_tag);
0131     pvas_token_ = consumes<edm::Association<VertexCollection>>(pvas_tag);
0132     use_pvasq_value_map_ = true;
0133   }
0134 }
0135
0136 void ParticleTransformerAK4TagInfoProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0137   // pfParticleTransformerAK4TagInfos
0138   edm::ParameterSetDescription desc;
0139   desc.add<double>("jet_radius", 0.4);
0140   desc.add<double>("min_candidate_pt", 0.95);
0141   desc.add<bool>("flip", false);
0142   desc.add<edm::InputTag>("vertices", edm::InputTag("offlinePrimaryVertices"));
0143   desc.add<edm::InputTag>("puppi_value_map", edm::InputTag("puppi"));
0144   desc.add<edm::InputTag>("secondary_vertices", edm::InputTag("inclusiveCandidateSecondaryVertices"));
0145   desc.add<edm::InputTag>("jets", edm::InputTag("ak4PFJetsCHS"));
0146   desc.add<edm::InputTag>("candidates", edm::InputTag("packedPFCandidates"));
0147   desc.add<edm::InputTag>("vertex_associator", edm::InputTag("primaryVertexAssociation", "original"));
0148   desc.add<bool>("fallback_puppi_weight", false);
0149   desc.add<bool>("fallback_vertex_association", false);
0150   desc.add<bool>("is_weighted_jet", false);
0151   desc.add<double>("min_jet_pt", 15.0);
0152   desc.add<double>("max_jet_eta", 2.5);
0153   descriptions.add("pfParticleTransformerAK4TagInfos", desc);
0154 }
0155
0156 void ParticleTransformerAK4TagInfoProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
0157   auto output_tag_infos = std::make_unique<ParticleTransformerAK4TagInfoCollection>();
0158   edm::Handle<edm::View<reco::Jet>> jets;
0159   iEvent.getByToken(jet_token_, jets);
0160
0161   edm::Handle<VertexCollection> vtxs;
0162   iEvent.getByToken(vtx_token_, vtxs);
0163   if (vtxs->empty()) {
0164     // produce empty TagInfos in case no primary vertex
0165     iEvent.put(std::move(output_tag_infos));
0166     return;  // exit event
0167   }
0168   // reference to primary vertex
0169   const auto& pv = vtxs->at(0);
0170
0171   edm::Handle<edm::View<reco::Candidate>> tracks;
0172   iEvent.getByToken(candidateToken_, tracks);
0173
0174   edm::Handle<SVCollection> svs;
0175   iEvent.getByToken(sv_token_, svs);
0176
0177   edm::Handle<edm::ValueMap<float>> puppi_value_map;
0178   if (use_puppi_value_map_) {
0179     iEvent.getByToken(puppi_value_map_token_, puppi_value_map);
0180   }
0181
0182   edm::Handle<edm::ValueMap<int>> pvasq_value_map;
0183   edm::Handle<edm::Association<VertexCollection>> pvas;
0184   if (use_pvasq_value_map_) {
0185     iEvent.getByToken(pvasq_value_map_token_, pvasq_value_map);
0186     iEvent.getByToken(pvas_token_, pvas);
0187   }
0188
0189   edm::ESHandle<TransientTrackBuilder> track_builder = iSetup.getHandle(track_builder_token_);
0190
0191   for (std::size_t jet_n = 0; jet_n < jets->size(); jet_n++) {
0192     // create data containing structure
0193     btagbtvdeep::ParticleTransformerAK4Features features;
0194
0195     // reco jet reference (use as much as possible)
0196     const auto& jet = jets->at(jet_n);
0197     if (jet.pt() < 15.0) {
0198       features.is_filled = false;
0199     }
0200     if (std::abs(jet.eta()) > 2.5) {
0201       features.is_filled = false;
0202     }
0203     // dynamical casting to pointers, null if not possible
0204     const auto* pf_jet = dynamic_cast<const reco::PFJet*>(&jet);
0205     const auto* pat_jet = dynamic_cast<const pat::Jet*>(&jet);
0206     edm::RefToBase<reco::Jet> jet_ref(jets, jet_n);
0207
0208     if (features.is_filled) {
0209       math::XYZVector jet_dir = jet.momentum().Unit();
0210       GlobalVector jet_ref_track_dir(jet.px(), jet.py(), jet.pz());
0211
0212       // copy which will be sorted
0213       auto svs_sorted = *svs;
0214       // sort by dxy
0215       std::sort(svs_sorted.begin(), svs_sorted.end(), [&pv](const auto& sva, const auto& svb) {
0216         return btagbtvdeep::sv_vertex_comparator(sva, svb, pv);
0217       });
0218       // fill features from secondary vertices
0219       for (const auto& sv : svs_sorted) {
0220         if (reco::deltaR2(sv, jet_dir) > (jet_radius_ * jet_radius_))
0221           continue;
0222         else {
0223           features.sv_features.emplace_back();
0224           // in C++17 could just get from emplace_back output
0225           auto& sv_features = features.sv_features.back();
0226           btagbtvdeep::svToFeatures(sv, pv, jet, sv_features, flip_);
0227         }
0228       }
0229
0230       // stuff required for dealing with pf candidates
0231
0232       std::vector<btagbtvdeep::SortingClass<size_t>> c_sorted, n_sorted;
0233
0234       // to cache the TrackInfo
0235       std::map<unsigned int, btagbtvdeep::TrackInfoBuilder> trackinfos;
0236
0237       // unsorted reference to sv
0238       const auto& svs_unsorted = *svs;
0239       // fill collection, from DeepTNtuples plus some styling
0240       for (unsigned int i = 0; i < jet.numberOfDaughters(); i++) {
0241         auto cand = jet.daughter(i);
0242         if (cand) {
0243           // candidates under 950MeV (configurable) are not considered
0244           // might change if we use also white-listing
0245           if (cand->pt() < min_candidate_pt_)
0246             continue;
0247           if (cand->charge() != 0) {
0248             auto& trackinfo = trackinfos.emplace(i, track_builder).first->second;
0249             trackinfo.buildTrackInfo(cand, jet_dir, jet_ref_track_dir, pv);
0250
0251             c_sorted.emplace_back(i,
0252                                   trackinfo.getTrackSip2dSig(),
0253                                   -btagbtvdeep::mindrsvpfcand(svs_unsorted, cand),
0254                                   cand->pt() / jet.pt());
0255           } else {
0256             n_sorted.emplace_back(i, -1, -btagbtvdeep::mindrsvpfcand(svs_unsorted, cand), cand->pt() / jet.pt());
0257           }
0258         }
0259       }
0260
0261       // sort collections (open the black-box if you please)
0262       std::sort(c_sorted.begin(), c_sorted.end(), btagbtvdeep::SortingClass<std::size_t>::compareByABCInv);
0263       std::sort(n_sorted.begin(), n_sorted.end(), btagbtvdeep::SortingClass<std::size_t>::compareByABCInv);
0264
0265       std::vector<size_t> c_sortedindices, n_sortedindices;
0266
0267       // this puts 0 everywhere and the right position in ind
0268       c_sortedindices = btagbtvdeep::invertSortingVector(c_sorted);
0269       n_sortedindices = btagbtvdeep::invertSortingVector(n_sorted);
0270
0271       // set right size to vectors
0272       features.c_pf_features.clear();
0273       features.c_pf_features.resize(c_sorted.size());
0274       features.n_pf_features.clear();
0275       features.n_pf_features.resize(n_sorted.size());
0276
0277       for (unsigned int i = 0; i < jet.numberOfDaughters(); i++) {
0278         // get pointer and check that is correct
0279         auto cand = dynamic_cast<const reco::Candidate*>(jet.daughter(i));
0280         if (!cand)
0281           continue;
0282         // candidates under 950MeV are not considered
0283         // might change if we use also white-listing
0284         if (cand->pt() < 0.95)
0285           continue;
0286
0287         auto packed_cand = dynamic_cast<const pat::PackedCandidate*>(cand);
0288         auto reco_cand = dynamic_cast<const reco::PFCandidate*>(cand);
0289
0290         // need some edm::Ptr or edm::Ref if reco candidates
0291         reco::PFCandidatePtr reco_ptr;
0292         if (pf_jet) {
0293           reco_ptr = pf_jet->getPFConstituent(i);
0294         } else if (pat_jet && reco_cand) {
0295           reco_ptr = pat_jet->getPFConstituent(i);
0296         }
0297
0298         reco::CandidatePtr cand_ptr;
0299         if (pat_jet) {
0300           cand_ptr = pat_jet->sourceCandidatePtr(i);
0301         }
0302
0303         //
0304         // Access puppi weight from ValueMap.
0305         //
0306         float puppiw = 1.0;  // Set to fallback value
0307
0308         if (reco_cand) {
0309           if (use_puppi_value_map_)
0310             puppiw = (*puppi_value_map)[reco_ptr];
0311           else if (!fallback_puppi_weight_) {
0312             throw edm::Exception(edm::errors::InvalidReference, "PUPPI value map missing")
0313                 << "use fallback_puppi_weight option to use " << puppiw << " for reco_cand as default";
0314           }
0315         } else if (packed_cand) {
0316           if (use_puppi_value_map_)
0317             puppiw = (*puppi_value_map)[cand_ptr];
0318           else if (!fallback_puppi_weight_) {
0319             throw edm::Exception(edm::errors::InvalidReference, "PUPPI value map missing")
0320                 << "use fallback_puppi_weight option to use " << puppiw << " for packed_cand as default";
0321           }
0322         } else {
0323           throw edm::Exception(edm::errors::InvalidReference)
0324               << "Cannot convert to either reco::PFCandidate or pat::PackedCandidate";
0325         }
0326
0327         float drminpfcandsv = btagbtvdeep::mindrsvpfcand(svs_unsorted, cand);
0328         float distminpfcandsv = 0;
0329         if (cand->charge() != 0) {
0330           // is charged candidate
0331           auto entry = c_sortedindices.at(i);
0332
0333           // get cached track info
0334           auto& trackinfo = trackinfos.at(i);
0335
0336           // get_ref to vector element
0337           auto& c_pf_features = features.c_pf_features.at(entry);
0338           // fill feature structure
0339           if (packed_cand) {
0340             if (packed_cand->hasTrackDetails()) {
0341               const reco::Track& PseudoTrack = packed_cand->pseudoTrack();
0342               reco::TransientTrack transientTrack;
0343               transientTrack = track_builder->build(PseudoTrack);
0344               distminpfcandsv = btagbtvdeep::mindistsvpfcand(svs_unsorted, transientTrack);
0345             }
0346
0347             btagbtvdeep::packedCandidateToFeatures(packed_cand,
0348                                                    jet,
0349                                                    trackinfo,
0350                                                    is_weighted_jet_,
0351                                                    drminpfcandsv,
0352                                                    static_cast<float>(jet_radius_),
0353                                                    puppiw,
0354                                                    c_pf_features,
0355                                                    flip_,
0356                                                    distminpfcandsv);
0357           } else if (reco_cand) {
0358             // get vertex association quality
0359             int pv_ass_quality = 0;  // fallback value
0360             if (use_pvasq_value_map_) {
0361               pv_ass_quality = (*pvasq_value_map)[reco_ptr];
0362             } else if (!fallback_vertex_association_) {
0363               throw edm::Exception(edm::errors::InvalidReference, "vertex association missing")
0364                   << "use fallback_vertex_association option to use" << pv_ass_quality
0365                   << "as default quality and closest dz PV as criteria";
0366             }
0367             // getting the PV as PackedCandidatesProducer
0368             // but using not the slimmed but original vertices
0369             auto ctrack = reco_cand->bestTrack();
0370             int pvi = -1;
0371             float dist = 1e99;
0372             for (size_t ii = 0; ii < vtxs->size(); ii++) {
0373               float dz = (ctrack) ? std::abs(ctrack->dz(((*vtxs)[ii]).position())) : 0;
0374               if (dz < dist) {
0375                 pvi = ii;
0376                 dist = dz;
0377               }
0378             }
0379             auto PV = reco::VertexRef(vtxs, pvi);
0380             if (use_pvasq_value_map_) {
0381               const reco::VertexRef& PV_orig = (*pvas)[reco_ptr];
0382               if (PV_orig.isNonnull())
0383                 PV = reco::VertexRef(vtxs, PV_orig.key());
0384             }
0385             btagbtvdeep::recoCandidateToFeatures(reco_cand,
0386                                                  jet,
0387                                                  trackinfo,
0388                                                  is_weighted_jet_,
0389                                                  drminpfcandsv,
0390                                                  static_cast<float>(jet_radius_),
0391                                                  puppiw,
0392                                                  pv_ass_quality,
0393                                                  PV,
0394                                                  c_pf_features,
0395                                                  flip_,
0396                                                  distminpfcandsv);
0397           }
0398         } else {
0399           // is neutral candidate
0400           auto entry = n_sortedindices.at(i);
0401           // get_ref to vector element
0402           auto& n_pf_features = features.n_pf_features.at(entry);
0403           // fill feature structure
0404           if (packed_cand) {
0405             btagbtvdeep::packedCandidateToFeatures(packed_cand,
0406                                                    jet,
0407                                                    is_weighted_jet_,
0408                                                    drminpfcandsv,
0409                                                    static_cast<float>(jet_radius_),
0410                                                    puppiw,
0411                                                    n_pf_features);
0412           } else if (reco_cand) {
0413             btagbtvdeep::recoCandidateToFeatures(
0414                 reco_cand, jet, is_weighted_jet_, drminpfcandsv, static_cast<float>(jet_radius_), puppiw, n_pf_features);
0415           }
0416         }
0417       }
0418     }
0419     output_tag_infos->emplace_back(features, jet_ref);
0420   }
0421   iEvent.put(std::move(output_tag_infos));
0422 }
0423
0424 //define this as a plug-in
0425 DEFINE_FWK_MODULE(ParticleTransformerAK4TagInfoProducer);