BPHNano/plugins/DiLeptonBuilder.cc

0001 /// original authors: RK18 team
0002 #include <algorithm>
0003 #include <limits>
0004 #include <map>
0005 #include <memory>
0006 #include <string>
0007 #include <vector>
0008
0009 #include "CommonTools/Statistics/interface/ChiSquaredProbability.h"
0010 #include "CommonTools/Utils/interface/StringCutObjectSelector.h"
0011 #include "DataFormats/BeamSpot/interface/BeamSpot.h"
0012 #include "DataFormats/Math/interface/deltaR.h"
0013 #include "DataFormats/PatCandidates/interface/CompositeCandidate.h"
0014 #include "FWCore/Framework/interface/Event.h"
0015 #include "FWCore/Framework/interface/global/EDProducer.h"
0016 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0017 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0018 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0019 #include "FWCore/Utilities/interface/InputTag.h"
0020 #include "KinVtxFitter.h"
0021 #include "TrackingTools/TransientTrack/interface/TransientTrack.h"
0022 #include "helper.h"
0023
0024 template <typename Lepton>
0025 class DiLeptonBuilder : public edm::global::EDProducer<> {
0026   // perhaps we need better structure here (begin run etc)
0027 public:
0028   typedef std::vector<Lepton> LeptonCollection;
0029   typedef std::vector<reco::TransientTrack> TransientTrackCollection;
0030
0031   explicit DiLeptonBuilder(const edm::ParameterSet &cfg)
0032       : l1_selection_{cfg.getParameter<std::string>("lep1Selection")},
0033         l2_selection_{cfg.getParameter<std::string>("lep2Selection")},
0034         pre_vtx_selection_{cfg.getParameter<std::string>("preVtxSelection")},
0035         post_vtx_selection_{cfg.getParameter<std::string>("postVtxSelection")},
0036         src_{consumes<LeptonCollection>(cfg.getParameter<edm::InputTag>("src"))},
0037         beamspot_{consumes<reco::BeamSpot>(cfg.getParameter<edm::InputTag>("beamSpot"))},
0038         ttracks_src_{consumes<TransientTrackCollection>(cfg.getParameter<edm::InputTag>("transientTracksSrc"))} {
0039     produces<pat::CompositeCandidateCollection>("SelectedDiLeptons");
0040   }
0041
0042   ~DiLeptonBuilder() override {}
0043
0044   void produce(edm::StreamID, edm::Event &, const edm::EventSetup &) const override;
0045
0046 private:
0047   const StringCutObjectSelector<Lepton> l1_selection_;                         // cut on leading lepton
0048   const StringCutObjectSelector<Lepton> l2_selection_;                         // cut on sub-leading lepton
0049   const StringCutObjectSelector<pat::CompositeCandidate> pre_vtx_selection_;   // cut on the di-lepton before the SV fit
0050   const StringCutObjectSelector<pat::CompositeCandidate> post_vtx_selection_;  // cut on the di-lepton after the SV fit
0051   const edm::EDGetTokenT<LeptonCollection> src_;
0052   const edm::EDGetTokenT<reco::BeamSpot> beamspot_;
0053   const edm::EDGetTokenT<TransientTrackCollection> ttracks_src_;
0054 };
0055
0056 template <typename Lepton>
0057 void DiLeptonBuilder<Lepton>::produce(edm::StreamID, edm::Event &evt, edm::EventSetup const &) const {
0058   // input
0059   edm::Handle<LeptonCollection> leptons;
0060   evt.getByToken(src_, leptons);
0061
0062   edm::Handle<reco::BeamSpot> beamspot;
0063   evt.getByToken(beamspot_, beamspot);
0064
0065   edm::Handle<TransientTrackCollection> ttracks;
0066   evt.getByToken(ttracks_src_, ttracks);
0067
0068   // output
0069   std::unique_ptr<pat::CompositeCandidateCollection> ret_value(new pat::CompositeCandidateCollection());
0070
0071   for (size_t l1_idx = 0; l1_idx < leptons->size(); ++l1_idx) {
0072     edm::Ptr<Lepton> l1_ptr(leptons, l1_idx);
0073     if (!l1_selection_(*l1_ptr))
0074       continue;
0075
0076     for (size_t l2_idx = l1_idx + 1; l2_idx < leptons->size(); ++l2_idx) {
0077       edm::Ptr<Lepton> l2_ptr(leptons, l2_idx);
0078       if (!l2_selection_(*l2_ptr))
0079         continue;
0080
0081       pat::CompositeCandidate lepton_pair;
0082       lepton_pair.setP4(l1_ptr->p4() + l2_ptr->p4());
0083       lepton_pair.setCharge(l1_ptr->charge() + l2_ptr->charge());
0084       lepton_pair.addUserFloat("lep_deltaR", reco::deltaR(*l1_ptr, *l2_ptr));
0085
0086       // Put the lepton passing the corresponding selection
0087       lepton_pair.addUserInt("l1_idx", l1_idx);
0088       lepton_pair.addUserInt("l2_idx", l2_idx);
0089
0090       // Use UserCands as they should not use memory but keep the Ptr itself
0091       lepton_pair.addUserCand("l1", l1_ptr);
0092       lepton_pair.addUserCand("l2", l2_ptr);
0093       if (!pre_vtx_selection_(lepton_pair))
0094         continue;  // before making the SV, cut on the info we have
0095
0096       KinVtxFitter fitter(
0097           {ttracks->at(l1_idx), ttracks->at(l2_idx)}, {l1_ptr->mass(), l2_ptr->mass()}, {bph::LEP_SIGMA, bph::LEP_SIGMA}
0098           // some small sigma for the particle mass
0099       );
0100       if (!fitter.success())
0101         continue;
0102       lepton_pair.setVertex(
0103           reco::Candidate::Point(fitter.fitted_vtx().x(), fitter.fitted_vtx().y(), fitter.fitted_vtx().z()));
0104       lepton_pair.addUserInt("sv_ok", fitter.success() ? 1 : 0);
0105       lepton_pair.addUserFloat("sv_chi2", fitter.chi2());
0106       lepton_pair.addUserFloat("sv_ndof", fitter.dof());  // float??
0107       lepton_pair.addUserFloat("sv_prob", fitter.prob());
0108       lepton_pair.addUserFloat("fitted_mass", fitter.success() ? fitter.fitted_candidate().mass() : -1);
0109       lepton_pair.addUserFloat(
0110           "fitted_massErr",
0111           fitter.success() ? sqrt(fitter.fitted_candidate().kinematicParametersError().matrix()(6, 6)) : -1);
0112       auto fit_p4 = fitter.fitted_p4();
0113       lepton_pair.addUserFloat("vtx_x", lepton_pair.vx());
0114       lepton_pair.addUserFloat("vtx_y", lepton_pair.vy());
0115       lepton_pair.addUserFloat("vtx_z", lepton_pair.vz());
0116       lepton_pair.addUserFloat("cos_theta_2D", bph::cos_theta_2D(fitter, *beamspot, lepton_pair.p4()));
0117       lepton_pair.addUserFloat("fitted_cos_theta_2D", bph::cos_theta_2D(fitter, *beamspot, fit_p4));
0118
0119       auto lxy = bph::l_xy(fitter, *beamspot);
0120       lepton_pair.addUserFloat("l_xy", lxy.value());
0121       lepton_pair.addUserFloat("l_xy_unc", lxy.error());
0122
0123       // cut on the SV info
0124       if (!post_vtx_selection_(lepton_pair))
0125         continue;
0126       ret_value->push_back(lepton_pair);
0127     }
0128   }
0129
0130   evt.put(std::move(ret_value), "SelectedDiLeptons");
0131 }
0132
0133 #include "DataFormats/PatCandidates/interface/Electron.h"
0134 #include "DataFormats/PatCandidates/interface/Muon.h"
0135 typedef DiLeptonBuilder<pat::Muon> DiMuonBuilder;
0136 typedef DiLeptonBuilder<pat::Electron> DiElectronBuilder;
0137
0138 #include "FWCore/Framework/interface/MakerMacros.h"
0139 DEFINE_FWK_MODULE(DiMuonBuilder);
0140 DEFINE_FWK_MODULE(DiElectronBuilder);