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