PatAlgos/plugins/LeptonTagInfoCollectionProducer.cc

0001 // -*- C++ -*-
0002 //
0003 // Package:    PhysicsTools/PatAlgos
0004 // Class:      LeptonTagInfoCollectionProducer
0005 //
0006 // Original Author:  Sergio Sanchez Cruz
0007 //         Created:  Mon, 15 May 2023 08:32:03 GMT
0008 //
0009 #include "FWCore/Framework/interface/Frameworkfwd.h"
0010 #include "FWCore/Framework/interface/stream/EDProducer.h"
0011
0012 #include "FWCore/Framework/interface/Event.h"
0013 #include "FWCore/Framework/interface/MakerMacros.h"
0014
0015 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0016 #include "FWCore/Utilities/interface/StreamID.h"
0017 #include "FWCore/Utilities/interface/ESGetToken.h"
0018
0019 #include "CommonTools/Utils/interface/StringObjectFunction.h"
0020 #include "DataFormats/BTauReco/interface/DeepBoostedJetFeatures.h"
0021 #include "DataFormats/Candidate/interface/VertexCompositePtrCandidate.h"
0022 #include "DataFormats/PatCandidates/interface/PackedCandidate.h"
0023 #include "DataFormats/PatCandidates/interface/Electron.h"
0024 #include "DataFormats/PatCandidates/interface/Muon.h"
0025 #include "DataFormats/PatCandidates/interface/Jet.h"
0026 #include "DataFormats/VertexReco/interface/Vertex.h"
0027 #include "DataFormats/Math/interface/deltaPhi.h"
0028 #include "DataFormats/Math/interface/deltaR.h"
0029
0030 #include "RecoVertex/VertexTools/interface/VertexDistance3D.h"
0031 #include "RecoVertex/VertexTools/interface/VertexDistanceXY.h"
0032 #include "RecoVertex/VertexPrimitives/interface/VertexState.h"
0033 #include "RecoVertex/VertexPrimitives/interface/ConvertToFromReco.h"
0034
0035 using namespace btagbtvdeep;
0036
0037 template <typename LeptonType>
0038 class LeptonTagInfoCollectionProducer : public edm::stream::EDProducer<> {
0039 public:
0040   explicit LeptonTagInfoCollectionProducer(const edm::ParameterSet& iConfig);
0041   ~LeptonTagInfoCollectionProducer() override {};
0042
0043   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0044
0045 private:
0046   using LeptonTagInfoCollection = DeepBoostedJetFeaturesCollection;
0047
0048   void produce(edm::Event& iEvent, const edm::EventSetup& iSetup) override;
0049
0050   void fill_lepton_features(const LeptonType&, DeepBoostedJetFeatures&);
0051   void fill_lepton_extfeatures(const edm::RefToBase<LeptonType>&, DeepBoostedJetFeatures&, edm::Event&);
0052   void fill_pf_features(const LeptonType&, DeepBoostedJetFeatures&);
0053   void fill_sv_features(const LeptonType&, DeepBoostedJetFeatures&);
0054
0055   template <typename VarType>
0056   using VarWithName = std::pair<std::string, StringObjectFunction<VarType, true>>;
0057   template <typename VarType>
0058   void parse_vars_into(const edm::ParameterSet& varsPSet, std::vector<std::unique_ptr<VarWithName<VarType>>>& vars) {
0059     for (const std::string& vname : varsPSet.getParameterNamesForType<std::string>()) {
0060       const std::string& func = varsPSet.getParameter<std::string>(vname);
0061       vars.push_back(std::make_unique<VarWithName<VarType>>(vname, StringObjectFunction<VarType, true>(func)));
0062     }
0063   }
0064
0065   template <typename VarType>
0066   using ExtVarWithName = std::pair<std::string, edm::EDGetTokenT<edm::ValueMap<VarType>>>;
0067   template <typename VarType>
0068   void parse_extvars_into(const edm::ParameterSet& varsPSet,
0069                           std::vector<std::unique_ptr<ExtVarWithName<VarType>>>& vars) {
0070     for (const std::string& vname : varsPSet.getParameterNamesForType<edm::InputTag>()) {
0071       vars.push_back(std::make_unique<ExtVarWithName<VarType>>(
0072           vname, consumes<edm::ValueMap<VarType>>(varsPSet.getParameter<edm::InputTag>(vname))));
0073     }
0074   }
0075
0076   edm::EDGetTokenT<edm::View<LeptonType>> src_token_;
0077   edm::EDGetTokenT<pat::PackedCandidateCollection> pf_token_;
0078   edm::EDGetTokenT<reco::VertexCompositePtrCandidateCollection> sv_token_;
0079   edm::EDGetTokenT<std::vector<reco::Vertex>> pv_token_;
0080
0081   edm::ParameterSet lepton_varsPSet_;
0082   edm::ParameterSet lepton_varsExtPSet_;
0083   edm::ParameterSet pf_varsPSet_;
0084   edm::ParameterSet sv_varsPSet_;
0085
0086   std::vector<std::unique_ptr<VarWithName<LeptonType>>> lepton_vars_;
0087   std::vector<std::unique_ptr<VarWithName<pat::PackedCandidate>>> pf_vars_;
0088   std::vector<std::unique_ptr<VarWithName<reco::VertexCompositePtrCandidate>>> sv_vars_;
0089   edm::Handle<reco::VertexCompositePtrCandidateCollection> svs_;
0090   edm::Handle<pat::PackedCandidateCollection> pfs_;
0091   edm::Handle<std::vector<reco::Vertex>> pvs_;
0092   std::vector<std::unique_ptr<ExtVarWithName<float>>> extLepton_vars_;
0093 };
0094
0095 template <typename LeptonType>
0096 LeptonTagInfoCollectionProducer<LeptonType>::LeptonTagInfoCollectionProducer(const edm::ParameterSet& iConfig)
0097     : src_token_(consumes<edm::View<LeptonType>>(iConfig.getParameter<edm::InputTag>("src"))),
0098       pf_token_(consumes<pat::PackedCandidateCollection>(iConfig.getParameter<edm::InputTag>("pfCandidates"))),
0099       sv_token_(consumes<reco::VertexCompositePtrCandidateCollection>(
0100           iConfig.getParameter<edm::InputTag>("secondary_vertices"))),
0101       pv_token_(consumes<std::vector<reco::Vertex>>(iConfig.getParameter<edm::InputTag>("pvSrc"))),
0102       lepton_varsPSet_(iConfig.getParameter<edm::ParameterSet>("leptonVars")),
0103       lepton_varsExtPSet_(iConfig.getParameter<edm::ParameterSet>("leptonVarsExt")),
0104       pf_varsPSet_(iConfig.getParameter<edm::ParameterSet>("pfVars")),
0105       sv_varsPSet_(iConfig.getParameter<edm::ParameterSet>("svVars")) {
0106   parse_vars_into(lepton_varsPSet_, lepton_vars_);
0107   parse_vars_into(pf_varsPSet_, pf_vars_);
0108   parse_vars_into(sv_varsPSet_, sv_vars_);
0109   parse_extvars_into(lepton_varsExtPSet_, extLepton_vars_);
0110
0111   produces<LeptonTagInfoCollection>();
0112 }
0113
0114 template <typename LeptonType>
0115 void LeptonTagInfoCollectionProducer<LeptonType>::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0116   edm::ParameterSetDescription desc;
0117
0118   desc.add<edm::InputTag>("src", edm::InputTag("slimmedMuons"));
0119   desc.add<edm::InputTag>("pfCandidates", edm::InputTag("packedPFCandidates"));
0120   desc.add<edm::InputTag>("secondary_vertices", edm::InputTag("slimmedSecondaryVertices"));
0121   desc.add<edm::InputTag>("pvSrc", edm::InputTag("offlineSlimmedPrimaryVertices"));
0122
0123   for (auto&& what : {"leptonVars", "pfVars", "svVars"}) {
0124     edm::ParameterSetDescription descNested;
0125     descNested.addWildcard<std::string>("*");
0126     desc.add<edm::ParameterSetDescription>(what, descNested);
0127   }
0128
0129   for (auto&& what : {"leptonVarsExt"}) {
0130     edm::ParameterSetDescription descNested;
0131     descNested.addWildcard<edm::InputTag>("*");
0132     desc.add<edm::ParameterSetDescription>(what, descNested);
0133   }
0134
0135   std::string modname;
0136   if (typeid(LeptonType) == typeid(pat::Muon))
0137     modname += "muon";
0138   else if (typeid(LeptonType) == typeid(pat::Electron))
0139     modname += "electron";
0140   modname += "TagInfos";
0141   descriptions.add(modname, desc);
0142 }
0143
0144 template <typename LeptonType>
0145 void LeptonTagInfoCollectionProducer<LeptonType>::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
0146   auto src = iEvent.getHandle(src_token_);
0147   iEvent.getByToken(sv_token_, svs_);
0148   iEvent.getByToken(pv_token_, pvs_);
0149   iEvent.getByToken(pf_token_, pfs_);
0150
0151   auto output_info = std::make_unique<LeptonTagInfoCollection>();
0152
0153   if (pvs_->empty()) {
0154     // produce empty TagInfos in case no primary vertex
0155     iEvent.put(std::move(output_info));
0156     return;
0157   }
0158
0159   for (size_t ilep = 0; ilep < src->size(); ilep++) {
0160     const auto& lep = (*src)[ilep];
0161     edm::RefToBase<LeptonType> lep_ref(src, ilep);
0162     DeepBoostedJetFeatures features;
0163     fill_lepton_features(lep, features);
0164     fill_lepton_extfeatures(lep_ref, features, iEvent);  // fixme
0165     fill_pf_features(lep, features);
0166     fill_sv_features(lep, features);
0167
0168     output_info->emplace_back(features);
0169   }
0170   iEvent.put(std::move(output_info));
0171 }
0172
0173 template <typename LeptonType>
0174 void LeptonTagInfoCollectionProducer<LeptonType>::fill_lepton_features(const LeptonType& lep,
0175                                                                        DeepBoostedJetFeatures& features) {
0176   for (auto& var : lepton_vars_) {
0177     features.add(var->first);
0178     features.reserve(var->first, 1);
0179     features.fill(var->first, var->second(lep));
0180   }
0181 }
0182
0183 template <typename LeptonType>
0184 void LeptonTagInfoCollectionProducer<LeptonType>::fill_lepton_extfeatures(const edm::RefToBase<LeptonType>& lep,
0185                                                                           DeepBoostedJetFeatures& features,
0186                                                                           edm::Event& iEvent) {
0187   for (auto& var : extLepton_vars_) {
0188     edm::Handle<edm::ValueMap<float>> vmap;
0189     iEvent.getByToken(var->second, vmap);
0190
0191     features.add(var->first);
0192     features.reserve(var->first, 1);
0193     features.fill(var->first, (*vmap)[lep]);
0194   }
0195 }
0196
0197 template <typename LeptonType>
0198 void LeptonTagInfoCollectionProducer<LeptonType>::fill_pf_features(const LeptonType& lep,
0199                                                                    DeepBoostedJetFeatures& features) {
0200   pat::PackedCandidateCollection pfcands;
0201   for (size_t ipf = 0; ipf < pfs_->size(); ++ipf) {
0202     if (reco::deltaR(pfs_->at(ipf), lep) < 0.4)
0203       pfcands.push_back(pfs_->at(ipf));
0204   }
0205
0206   for (auto& var : pf_vars_) {
0207     features.add(var->first);
0208     features.reserve(var->first, pfcands.size());
0209     for (const auto& cand : pfcands) {
0210       features.fill(var->first, var->second(cand));
0211     }
0212   }
0213
0214   // afaik these need to be hardcoded because I cannot put userFloats to pat::packedCandidates
0215   features.add("PF_phi_rel");
0216   features.reserve("PF_phi_rel", pfcands.size());
0217   features.add("PF_eta_rel");
0218   features.reserve("PF_eta_rel", pfcands.size());
0219   features.add("PF_dR_lep");
0220   features.reserve("PF_dR_lep", pfcands.size());
0221   features.add("PF_pt_rel_log");
0222   features.reserve("PF_pt_rel_log", pfcands.size());
0223
0224   for (const auto& cand : pfcands) {
0225     features.fill("PF_phi_rel", reco::deltaPhi(lep.phi(), cand.phi()));
0226     features.fill("PF_eta_rel", lep.eta() - cand.eta());
0227     features.fill("PF_dR_lep", reco::deltaR(lep, cand));
0228     features.fill("PF_pt_rel_log", log(cand.pt() / lep.pt()));
0229   }
0230 }
0231
0232 template <typename LeptonType>
0233 void LeptonTagInfoCollectionProducer<LeptonType>::fill_sv_features(const LeptonType& lep,
0234                                                                    DeepBoostedJetFeatures& features) {
0235   reco::VertexCompositePtrCandidateCollection selectedSVs;
0236   for (size_t isv = 0; isv < svs_->size(); ++isv) {
0237     if (reco::deltaR(lep, svs_->at(isv)) < 0.4) {
0238       selectedSVs.push_back(svs_->at(isv));
0239     }
0240   }
0241
0242   for (auto& var : sv_vars_) {
0243     features.add(var->first);
0244     features.reserve(var->first, selectedSVs.size());
0245     for (auto& sv : selectedSVs)
0246       features.fill(var->first, var->second(sv));
0247   }
0248
0249   // afaik these need to be hardcoded
0250   const auto& PV0 = pvs_->front();
0251   VertexDistance3D vdist;
0252   VertexDistanceXY vdistXY;
0253
0254   features.add("SV_dlenSig");
0255   features.reserve("SV_dlenSig", selectedSVs.size());
0256   features.add("SV_dxy");
0257   features.reserve("SV_dxy", selectedSVs.size());
0258   features.add("SV_eta_rel");
0259   features.reserve("SV_eta_rel", selectedSVs.size());
0260   features.add("SV_phi_rel");
0261   features.reserve("SV_phi_rel", selectedSVs.size());
0262   features.add("SV_dR_lep");
0263   features.reserve("SV_dR_lep", selectedSVs.size());
0264   features.add("SV_pt_rel");
0265   features.reserve("SV_pt_rel", selectedSVs.size());
0266   features.add("SV_cospAngle");
0267   features.reserve("SV_cospAngle", selectedSVs.size());
0268   features.add("SV_d3d");
0269   features.reserve("SV_d3d", selectedSVs.size());
0270
0271   for (auto& sv : selectedSVs) {
0272     Measurement1D dl =
0273         vdist.distance(PV0, VertexState(RecoVertex::convertPos(sv.position()), RecoVertex::convertError(sv.error())));
0274     features.fill("SV_d3d", dl.value());
0275     features.fill("SV_dlenSig", dl.significance());
0276     Measurement1D d2d =
0277         vdistXY.distance(PV0, VertexState(RecoVertex::convertPos(sv.position()), RecoVertex::convertError(sv.error())));
0278     features.fill("SV_dxy", d2d.value());
0279     features.fill("SV_phi_rel", reco::deltaPhi(lep.phi(), sv.phi()));
0280     features.fill("SV_eta_rel", lep.eta() - sv.eta());
0281     features.fill("SV_dR_lep", reco::deltaR(sv, lep));
0282     features.fill("SV_pt_rel", sv.pt() / lep.pt());
0283     double dx = (PV0.x() - sv.vx()), dy = (PV0.y() - sv.vy()), dz = (PV0.z() - sv.vz());
0284     double pdotv = (dx * sv.px() + dy * sv.py() + dz * sv.pz()) / sv.p() / sqrt(dx * dx + dy * dy + dz * dz);
0285     features.fill("SV_cospAngle", pdotv);
0286   }
0287 }
0288
0289 typedef LeptonTagInfoCollectionProducer<pat::Muon> MuonTagInfoCollectionProducer;
0290 typedef LeptonTagInfoCollectionProducer<pat::Electron> ElectronTagInfoCollectionProducer;
0291
0292 DEFINE_FWK_MODULE(MuonTagInfoCollectionProducer);
0293 DEFINE_FWK_MODULE(ElectronTagInfoCollectionProducer);